Repository: Toblerity/Shapely
Pull request #1177 (2021-08-03 08:39:28)

Refactor STRtree implementation to store geoms/idx sequences instead of reverse mapping
This PR takes the implementation details from https://github.com/Toblerity/Shapely/pull/1064 (i.e. using  `geoms` and `idxs` sequences (lists) to store the data with which the GEOSSTRtree gets created, instead of using a mapping), but cleanly on top of the new interface (`query_geoms`/`query_items`) from https://github.com/Toblerity/Shapely/pull/1112, without all the discussion on the interface from #1064.

It doesn't change any user facing API now after #1112, but only refactors the implementation (how to store the items). As mentioned in https://github.com/Toblerity/Shapely/pull/1112#issuecomment-875108728, this version is closer to what we would do in the shapely-2.0 branch on top of the C-based STRtree (which I started in https://github.com/Toblerity/Shapely/pull/1161). IMO it is useful to already do this now (it will also make it a bit easier to have the master and shapely-2.0 branches next to each other for a while).

coveralls: 
[![Coverage Status](https://coveralls.io/builds/41857513/badge)](https://coveralls.io/builds/41857513)

Coverage increased (+0.02%) to 85.31% when pulling **6a10369ea91edc1ee407b843d2e46f49df94eb25 on jorisvandenbossche:strtree-refactor** into **9a761739aea19823fcead4391ef043edafa74db6 on Toblerity:master**.

jorisvandenbossche: cc @sgillies what do you think about this PR?
sgillies: I'll merge and then we could have another pre-release on Monday.

Patch:
diff --git a/shapely/coords.py b/shapely/coords.py
index df673b0..ce526ef 100644
--- a/shapely/coords.py
+++ b/shapely/coords.py
@@ -116,7 +116,10 @@ class CoordinateSequence:
     def _ctypes(self):
         self._update()
         has_z = self._ndim == 3
-        n = self._ndim
+        n = self._ndim or 0
+        if n == 0:
+            # ignore with NumPy 1.21 __array_interface__
+            raise AttributeError("empty geometry sequence")
         m = self.__len__()
         array_type = c_double * (m * n)
         data = array_type()
diff --git a/shapely/geometry/base.py b/shapely/geometry/base.py
index 1ad0d07..2a5eacb 100644
--- a/shapely/geometry/base.py
+++ b/shapely/geometry/base.py
@@ -849,8 +849,8 @@ class BaseMultipartGeometry(BaseGeometry):
     @property
     def __array_interface__(self):
         """Provide the Numpy array protocol."""
-        raise NotImplementedError("Multi-part geometries do not themselves "
-                                  "provide the array interface")
+        raise AttributeError("Multi-part geometries do not themselves "
+                             "provide the array interface")
 
     def _get_coords(self):
         raise NotImplementedError("Sub-geometries may have coordinate "
diff --git a/shapely/geometry/polygon.py b/shapely/geometry/polygon.py
index 4386eb4..fabb6f2 100644
--- a/shapely/geometry/polygon.py
+++ b/shapely/geometry/polygon.py
@@ -317,7 +317,7 @@ class Polygon(BaseGeometry):
 
     @property
     def __array_interface__(self):
-        raise NotImplementedError(
+        raise AttributeError(
         "A polygon does not itself provide the array interface. Its rings do.")
 
     def _get_coords(self):
diff --git a/shapely/speedups/_speedups.pyx b/shapely/speedups/_speedups.pyx
index f5b6dc2..0b49959 100644
--- a/shapely/speedups/_speedups.pyx
+++ b/shapely/speedups/_speedups.pyx
@@ -421,7 +421,10 @@ def coordseq_ctypes(self):
     cdef GEOSCoordSequence *cs
     cdef double *data_p
     self._update()
-    n = self._ndim
+    n = self._ndim or 0
+    if n == 0:
+        # ignore with NumPy 1.21 __array_interface__
+        raise AttributeError("empty geometry sequence")
     m = self.__len__()
     array_type = ctypes.c_double * (m * n)
     data = array_type()
diff --git a/shapely/strtree.py b/shapely/strtree.py
index 30c79a5..a345058 100644
--- a/shapely/strtree.py
+++ b/shapely/strtree.py
@@ -104,44 +104,34 @@ class STRtree:
             ShapelyDeprecationWarning,
             stacklevel=2,
         )
-        self._tree = None
         self.node_capacity = node_capacity
-        self._rev = {
-            item: geom
-            for geom, item in self._iterinitdata(geoms, items)
-            if not geom.is_empty
-        }
-        if self._rev:
-            self._init_tree(self._rev.items())
-
-    def _iterinitdata(
-        self,
-        geoms: Iterable[BaseGeometry], items: Optional[Iterable[BaseGeometry]],
-    ) -> Iterator[Tuple[BaseGeometry, Any]]:
-        if items is not None:
-            for geom, item in zip(geoms, items):
-                if isinstance(geom, BaseGeometry):
-                    yield (geom, item)
+
+        # 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
+
+        # initialize GEOS STRtree
+        self._tree = lgeos.GEOSSTRtree_create(self.node_capacity)
+        i = 0
+        for idx, geom in zip(self._idxs, self._geoms):
+            # 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):
+        if self._has_custom_items:
+            return STRtree, (self._geoms, self._items)
         else:
-            for enum_idx, geom in enumerate(geoms):
-                if isinstance(geom, BaseGeometry):
-                    yield (geom, enum_idx)
-
-    def _init_tree(self, rev_initdata: ItemsView[Any, BaseGeometry]):
-        if rev_initdata:
-            self._tree = lgeos.GEOSSTRtree_create(self.node_capacity)
-            for item, geom in rev_initdata:
-                lgeos.GEOSSTRtree_insert(self._tree, geom._geom, ctypes.py_object(item))
-
-    def __getstate__(self):
-        state = self.__dict__.copy()
-        del state["_tree"]
-        return state
-
-    def __setstate__(self, state):
-        self.__dict__.update(state)
-        if self._rev:
-            self._init_tree(self._rev.items())
+            return STRtree, (self._geoms, )
 
     def __del__(self):
         if self._tree is not None:
@@ -152,6 +142,19 @@ class STRtree:
 
             self._tree = None
 
+    def _query(self, geom):
+        if self._n_geoms == 0:
+            return []
+
+        result = []
+
+        def callback(item, userdata):
+            idx = ctypes.cast(item, ctypes.py_object).value
+            result.append(idx)
+
+        lgeos.GEOSSTRtree_query(self._tree, geom._geom, lgeos.GEOSQueryCallback(callback), None)
+        return result
+
     def query_items(self, geom: BaseGeometry) -> Sequence[Any]:
         """Query for nodes which intersect the geom's envelope to get
         stored items.
@@ -197,19 +200,11 @@ class STRtree:
         ['POINT (2 2)']
 
         """
-        if self._tree is None or not self._rev:
-            return []
-
-        result = []
-
-        def callback(item, userdata):
-            idx = ctypes.cast(item, ctypes.py_object).value
-            result.append(idx)
-
-        lgeos.GEOSSTRtree_query(
-            self._tree, geom._geom, lgeos.GEOSQueryCallback(callback), None
-        )
-        return result
+        result = self._query(geom)
+        if self._has_custom_items:
+            return [self._items[i] for i in result]
+        else:
+            return result
 
     def query_geoms(self, geom: BaseGeometry) -> Sequence[BaseGeometry]:
         """Query for nodes which intersect the geom's envelope to get
@@ -225,8 +220,8 @@ class STRtree:
         An array or list of geometry objects.
 
         """
-        items = self.query_items(geom)
-        return [self._rev[idx] for idx in items]
+        result = self._query(geom)
+        return [self._geoms[i] for i in result]
 
     def query(self, geom: BaseGeometry) -> Sequence[BaseGeometry]:
         """Query for nodes which intersect the geom's envelope to get
@@ -247,6 +242,34 @@ class STRtree:
         """
         return self.query_geoms(geom)
 
+    def _nearest(self, geom, exclusive):
+        envelope = geom.envelope
+
+        def callback(item1, item2, distance, userdata):
+            try:
+                callback_userdata = ctypes.cast(userdata, ctypes.py_object).value
+                idx = ctypes.cast(item1, ctypes.py_object).value
+                geom2 = ctypes.cast(item2, ctypes.py_object).value
+                dist = ctypes.cast(distance, ctypes.POINTER(ctypes.c_double))
+                if callback_userdata["exclusive"] and self._geoms[idx].equals(geom2):
+                    dist[0] = sys.float_info.max
+                else:
+                    lgeos.GEOSDistance(self._geoms[idx]._geom, geom2._geom, dist)
+                
+                return 1
+            except Exception:
+                log.exception("Caught exception")
+                return 0
+
+        item = lgeos.GEOSSTRtree_nearest_generic(
+            self._tree,
+            ctypes.py_object(geom),
+            envelope._geom,
+            lgeos.GEOSDistanceCallback(callback),
+            ctypes.py_object({"exclusive": exclusive}),
+        )
+        return ctypes.cast(item, ctypes.py_object).value
+
     def nearest_item(
         self, geom: BaseGeometry, exclusive: bool = False
     ) -> Union[Any, None]:
@@ -285,35 +308,14 @@ class STRtree:
         'POINT (0 0)'
 
         """
-        if self._tree is None or not self._rev:
+        if self._n_geoms == 0:
             return None
 
-        envelope = geom.envelope
-
-        def callback(item1, item2, distance, userdata):
-            try:
-                callback_userdata = ctypes.cast(userdata, ctypes.py_object).value
-                idx = ctypes.cast(item1, ctypes.py_object).value
-                geom2 = ctypes.cast(item2, ctypes.py_object).value
-                dist = ctypes.cast(distance, ctypes.POINTER(ctypes.c_double))
-                if callback_userdata["exclusive"] and self._rev[idx].equals(geom2):
-                    dist[0] = sys.float_info.max
-                else:
-                    lgeos.GEOSDistance(self._rev[idx]._geom, geom2._geom, dist)
-                return 1
-            except Exception:
-                log.exception("Caught exception")
-                return 0
-
-        item = lgeos.GEOSSTRtree_nearest_generic(
-            self._tree,
-            ctypes.py_object(geom),
-            envelope._geom,
-            lgeos.GEOSDistanceCallback(callback),
-            ctypes.py_object({"exclusive": exclusive}),
-        )
-        result = ctypes.cast(item, ctypes.py_object).value
-        return result
+        result = self._nearest(geom, exclusive)
+        if self._has_custom_items:
+            return self._items[result]
+        else:
+            return result
 
     def nearest_geom(
         self, geom: BaseGeometry, exclusive: bool = False
@@ -337,11 +339,8 @@ class STRtree:
         version 2.0.
 
         """
-        item = self.nearest_item(geom, exclusive=exclusive)
-        if item is None:
-            return None
-        else:
-            return self._rev[item]
+        result = self._nearest(geom, exclusive)
+        return self._geoms[result]
 
     def nearest(
         self, geom: BaseGeometry, exclusive: bool = False


Test patch:
diff --git a/tests/test_collection.py b/tests/test_collection.py
index 09b5c11..5ba176b 100644
--- a/tests/test_collection.py
+++ b/tests/test_collection.py
@@ -114,3 +114,14 @@ def test_len_deprecated(geometrycollection_geojson):
     geom = shape(geometrycollection_geojson)
     with pytest.warns(ShapelyDeprecationWarning, match="__len__"):
         assert len(geom) == 2
+
+
+@shapely20_deprecated
+@pytest.mark.filterwarnings("error:An exception was ignored")  # NumPy 1.21
+def test_numpy_object_array():
+    np = pytest.importorskip("numpy")
+
+    geom = GeometryCollection([LineString([(0, 0), (1, 1)])])
+    ar = np.empty(1, object)
+    ar[:] = [geom]
+    assert ar[0] == geom
diff --git a/tests/test_linestring.py b/tests/test_linestring.py
index 2b6e7f9..a452ff7 100644
--- a/tests/test_linestring.py
+++ b/tests/test_linestring.py
@@ -87,6 +87,7 @@ def test_from_numpy():
     assert line.coords[:] == [(1.0, 2.0), (3.0, 4.0)]
 
 
+@pytest.mark.filterwarnings("error:An exception was ignored")  # NumPy 1.21
 def test_numpy_empty_linestring_coords():
     np = pytest.importorskip("numpy")
 
@@ -97,6 +98,17 @@ def test_numpy_empty_linestring_coords():
     assert la.shape == (0,)
 
 
+@shapely20_deprecated
+@pytest.mark.filterwarnings("error:An exception was ignored")  # NumPy 1.21
+def test_numpy_object_array():
+    np = pytest.importorskip("numpy")
+
+    geom = LineString([(0.0, 0.0), (0.0, 1.0)])
+    ar = np.empty(1, object)
+    ar[:] = [geom]
+    assert ar[0] == geom
+
+
 def test_from_invalid_dim():
     with pytest.raises(ValueError, match="at least 2 coordinate tuples"):
         LineString([(1, 2)])
diff --git a/tests/test_multilinestring.py b/tests/test_multilinestring.py
index bf79ddf..c6475dc 100644
--- a/tests/test_multilinestring.py
+++ b/tests/test_multilinestring.py
@@ -117,3 +117,14 @@ def test_getitem_deprecated():
     geom = MultiLineString([[[5.0, 6.0], [7.0, 8.0]]])
     with pytest.warns(ShapelyDeprecationWarning, match="__getitem__"):
         part = geom[0]
+
+
+@shapely20_deprecated
+@pytest.mark.filterwarnings("error:An exception was ignored")  # NumPy 1.21
+def test_numpy_object_array():
+    np = pytest.importorskip("numpy")
+
+    geom = MultiLineString([[[5.0, 6.0], [7.0, 8.0]]])
+    ar = np.empty(1, object)
+    ar[:] = [geom]
+    assert ar[0] == geom
diff --git a/tests/test_multipoint.py b/tests/test_multipoint.py
index afbffaf..a3e2914 100644
--- a/tests/test_multipoint.py
+++ b/tests/test_multipoint.py
@@ -129,6 +129,17 @@ def test_multipoint_array_interface_numpy_deprecated():
         np.array(geom)
 
 
+@shapely20_deprecated
+@pytest.mark.filterwarnings("error:An exception was ignored")  # NumPy 1.21
+def test_numpy_object_array():
+    np = pytest.importorskip("numpy")
+
+    geom = MultiPoint(((1.0, 2.0), (3.0, 4.0)))
+    ar = np.empty(1, object)
+    ar[:] = [geom]
+    assert ar[0] == geom
+
+
 def test_iteration_deprecated():
     geom = MultiPoint([[5.0, 6.0], [7.0, 8.0]])
     with pytest.warns(ShapelyDeprecationWarning, match="Iteration"):
diff --git a/tests/test_multipolygon.py b/tests/test_multipolygon.py
index c3982a8..878211a 100644
--- a/tests/test_multipolygon.py
+++ b/tests/test_multipolygon.py
@@ -110,3 +110,16 @@ def test_getitem_deprecated():
           [((0.25, 0.25), (0.25, 0.5), (0.5, 0.5), (0.5, 0.25))])])
     with pytest.warns(ShapelyDeprecationWarning, match="__getitem__"):
         part = geom[0]
+
+
+@shapely20_deprecated
+@pytest.mark.filterwarnings("error:An exception was ignored")  # NumPy 1.21
+def test_numpy_object_array():
+    np = pytest.importorskip("numpy")
+
+    geom = MultiPolygon(
+        [(((0.0, 0.0), (0.0, 1.0), (1.0, 1.0), (1.0, 0.0)),
+          [((0.25, 0.25), (0.25, 0.5), (0.5, 0.5), (0.5, 0.25))])])
+    ar = np.empty(1, object)
+    ar[:] = [geom]
+    assert ar[0] == geom
diff --git a/tests/test_point.py b/tests/test_point.py
index 264118c..4a80564 100644
--- a/tests/test_point.py
+++ b/tests/test_point.py
@@ -262,6 +262,7 @@ def test_point_array_interface_numpy_deprecated():
         np.array(p)
 
 
+@pytest.mark.filterwarnings("error:An exception was ignored")  # NumPy 1.21
 def test_numpy_empty_point_coords():
     np = pytest.importorskip("numpy")
 
@@ -270,3 +271,14 @@ def test_numpy_empty_point_coords():
     # Access the coords
     a = np.asarray(pe.coords)
     assert a.shape == (0,)
+
+
+@shapely20_deprecated
+@pytest.mark.filterwarnings("error:An exception was ignored")  # NumPy 1.21
+def test_numpy_object_array():
+    np = pytest.importorskip("numpy")
+
+    geom = Point(3.0, 4.0)
+    ar = np.empty(1, object)
+    ar[:] = [geom]
+    assert ar[0] == geom
diff --git a/tests/test_polygon.py b/tests/test_polygon.py
index 863b623..62b2f96 100644
--- a/tests/test_polygon.py
+++ b/tests/test_polygon.py
@@ -109,6 +109,7 @@ def test_numpy_linearring_coords():
     assert_array_equal(ra, expected)
 
 
+@pytest.mark.filterwarnings("error:An exception was ignored")  # NumPy 1.21
 def test_numpy_empty_linearring_coords():
     np = pytest.importorskip("numpy")
 
@@ -116,6 +117,16 @@ def test_numpy_empty_linearring_coords():
     assert np.asarray(ring.coords).shape == (0,)
 
 
+@pytest.mark.filterwarnings("error:An exception was ignored")  # NumPy 1.21
+def test_numpy_object_array():
+    np = pytest.importorskip("numpy")
+
+    geom = Polygon([(0.0, 0.0), (0.0, 1.0), (1.0, 1.0)])
+    ar = np.empty(1, object)
+    ar[:] = [geom]
+    assert ar[0] == geom
+
+
 def test_polygon_from_coordinate_sequence():
     coords = [(0.0, 0.0), (0.0, 1.0), (1.0, 1.0), (0.0, 0.0)]