fix returned index

jorisvandenbossche · jorisvandenbossche · commit 3bc928fc36ed · 2021-07-17T12:27:39.000+02:00
diff --git a/shapely/strtree.py b/shapely/strtree.py
@@ -116,28 +116,34 @@ def __init__(
             # unzip the list of tuples
             geoms, items = list(zip(*geoms))
 
-        has_custom_items = items is not None
-        if not has_custom_items:
-            items = list(range(len(geoms)))
+        # Keep references to geoms
+        self._geoms = list(geoms)
+        # Default enumeration index to store in the tree
+        self._idxs = list(range(len(self._geoms)))
 
+        # handle items
+        self._has_custom_items = items is not None
+        if not self._has_custom_items:
+            items = self._idxs
         self._items = items
-        self._has_custom_items = has_custom_items
-        # filter empty geometries out of the input
-        geoms = [geom for geom in geoms if not geom.is_empty]
-        self._n_geoms = len(geoms)
-        self._init_tree(geoms)
 
-        # Keep references to geoms.
-        self._geoms = list(geoms)
+        self._init_tree(geoms)
 
     def _init_tree(self, geoms):
         self._tree = lgeos.GEOSSTRtree_create(self.node_capacity)
-        self._idxs = list(range(len(geoms)))
+        i = 0
         for idx, geom in zip(self._idxs, geoms):
-            lgeos.GEOSSTRtree_insert(self._tree, geom._geom, ctypes.py_object(idx))
+            # filter empty geometries out of the input
+            if geom is not None and not geom.is_empty:
+                lgeos.GEOSSTRtree_insert(self._tree, geom._geom, ctypes.py_object(idx))
+                i += 1
+        self._n_geoms = i
 
     def __reduce__(self):
-        return STRtree, (self._geoms, self._items)
+        if self._has_custom_items:
+            return STRtree, (self._geoms, self._items)
+        else:
+            return STRtree, (self._geoms, )
 
     def __del__(self):
         if self._tree is not None:
diff --git a/tests/test_strtree.py b/tests/test_strtree.py
@@ -7,14 +7,19 @@
 import pytest
 
 from shapely.errors import ShapelyDeprecationWarning
-from shapely.geometry import Point, Polygon
+from shapely.geometry import Point, Polygon, box
 from shapely.geos import geos_version
 from shapely import strtree
 from shapely.strtree import STRtree
+from shapely import wkt
 
 from .conftest import requires_geos_342
 
 
+point = Point(2, 3)
+empty = wkt.loads("GEOMETRYCOLLECTION EMPTY")
+
+
 @requires_geos_342
 @pytest.mark.parametrize("geoms", [[Point(i, i) for i in range(10)]])
 @pytest.mark.parametrize(
@@ -42,6 +47,37 @@ def test_query_enumeration_idx(geoms, query_geom, expected):
     assert sorted(results) == sorted(expected)
 
 
+@requires_geos_342
+@pytest.mark.parametrize("geoms", [[Point(i, i) for i in range(5)]])
+@pytest.mark.parametrize("items", [None, list(range(1, 6)), list("abcde")])
+@pytest.mark.parametrize(
+    "query_geom,expected",
+    [(Point(2, 2).buffer(0.99), [2]), (Point(2, 2).buffer(1.0), [1, 2, 3])],
+)
+def test_query_items(geoms, items, query_geom, expected):
+    """Store enumeration idx"""
+    with pytest.warns(ShapelyDeprecationWarning):
+        tree = STRtree(geoms, items)
+    results = tree.query_items(query_geom)
+    expected = [items[idx] for idx in expected] if items is not None else expected
+    assert sorted(results) == sorted(expected)
+
+
+@pytest.mark.parametrize(
+    "tree_geometry, geometry,expected",
+    [
+        ([point], box(0, 0, 10, 10), [0]),
+        # None/empty is ignored in the tree, but the index of the valid geometry
+        # should be retained.
+        ([None, point], box(0, 0, 10, 10), [1]),
+        ([None, empty, point], box(0, 0, 10, 10), [2]),
+    ],
+)
+def test_query_items_with_empty(tree_geometry, geometry, expected):
+    with pytest.warns(ShapelyDeprecationWarning):
+        tree = STRtree(tree_geometry)
+    assert tree.query_items(geometry) == expected
+
 @requires_geos_342
 def test_insert_empty_geometry():
     """
