A general purpose STRtree which can store any Python object

shapely/strtree.py

@@ -19,11 +19,38 @@
 """
 
 import ctypes
+from functools import wraps
+import logging
 from warnings import warn
 
 from shapely.errors import ShapelyDeprecationWarning
 from shapely.geos import lgeos
 
+log = logging.getLogger(__name__)
+
+
+def nearest_callback(func):
+    @wraps(func)
+    def wrapper(arg1, arg2, arg3, arg4):
+        value = ctypes.cast(arg1, ctypes.py_object).value
+        geom = ctypes.cast(arg2, ctypes.py_object).value
+        dist = ctypes.cast(arg3, ctypes.POINTER(ctypes.c_double))
+        try:
+            dist.contents.value = func(value, geom)
+            return 1
+        except Exception:
+            log.exception()
+            return 0
+    return wrapper
+
+
+def query_callback(func):
+    @wraps(func)
+    def wrapper(arg1, arg2):
+        value = ctypes.cast(arg1, ctypes.py_object).value
+        func(value)
+    return wrapper
+
 
 class STRtree:
     """
@@ -75,49 +102,70 @@ class STRtree:
     None
     """
 
-    def __init__(self, geoms):
+    def __init__(self, initdata=None):
+        """Create a new STRtree.
+
+        Parameters
+        ----------
+        initdata : iterable
+            An iterable sequence of single geometry objects or (geom,
+            value) tuples.
+
+        Notes
+        -----
+        Items from initdata will be stored in two lists.
+
+        """
         warn(
             "STRtree will be completely changed in 2.0.0. The exact API is not yet decided, but will be documented before 1.8.0",
             ShapelyDeprecationWarning,
             stacklevel=2,
         )
-        # 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_handle(geoms)
-
-        # Keep references to geoms.
-        self._geoms = list(geoms)
-
-    def _init_tree_handle(self, geoms):
-        node_capacity = 10
-        self._tree_handle = lgeos.GEOSSTRtree_create(node_capacity)
-        for geom in geoms:
-            lgeos.GEOSSTRtree_insert(self._tree_handle, geom._geom, ctypes.py_object(geom))
+        self._initdata = None
+        self._tree = None
+        if initdata:
+            self._initdata = list(self._iteritems(initdata))
+            self._init_tree(self._initdata)
+
+    def _iteritems(self, initdata):
+        for obj in initdata:
+            if not isinstance(obj, tuple):
+                geom = obj
+                value = obj
+            else:
+                value, geom = obj
+            if not geom.is_empty:
+                yield value, geom
+
+    def _init_tree(self, initdata):
+        if initdata:
+            node_capacity = 10
+            self._tree = lgeos.GEOSSTRtree_create(node_capacity)
+            for value, geom in self._iteritems(initdata):
+                lgeos.GEOSSTRtree_insert(
+                    self._tree, geom._geom, ctypes.py_object(value)
+                )
 
     def __getstate__(self):
         state = self.__dict__.copy()
-        del state["_tree_handle"]
+        del state["_tree"]
         return state
 
     def __setstate__(self, state):
         self.__dict__.update(state)
-        self._init_tree_handle(self._geoms)
+        self._init_tree(self._initdata)
 
     def __del__(self):
-        if self._tree_handle is not None:
+        if self._tree is not None:
             try:
-                lgeos.GEOSSTRtree_destroy(self._tree_handle)
-            except AttributeError:
-                pass  # lgeos might be empty on shutdown.
-
-            self._tree_handle = None
+                lgeos.GEOSSTRtree_destroy(self._tree)
+            except AttributeError:  # lgeos might be empty on shutdown.
+                pass
+            self._tree = None
 
     def query(self, geom):
         """
-        Search the index for geometry objects whose extents
-        intersect the extent of the given object.
+        Search the tree for nodes which intersect geom's envelope
 
         Parameters
         ----------
@@ -126,18 +174,17 @@ def query(self, geom):
 
         Returns
         -------
-        list of geometry objects
-            All the geometry objects in the index whose extents
-            intersect the extent of `geom`.
+        list
+            A list of the values stored at the found nodes. The list
+            will be empty if the tree is empty.
 
         Note
         ----
-        A geometry object's "extent" is its the minimum xy bounding
+        A geometry object's envelope is its the minimum xy bounding
         rectangle.
 
         Examples
         --------
-
         A buffer around a point can be used to control the extent
         of the query.
 
@@ -160,45 +207,53 @@ def query(self, geom):
         >>> [o.wkt for o in tree.query(query_geom) if o.intersects(query_geom)]
         ['POINT (2 2)']
 
-        To get the original indices of the returned objects, create an
-        auxiliary dictionary. But use the geometry *ids* as keys since
-        the shapely geometry objects themselves are not hashable.
-
-        >>> index_by_id = dict((id(pt), i) for i, pt in enumerate(points))
-        >>> [(index_by_id[id(pt)], pt.wkt) for pt in tree.query(Point(2,2).buffer(1.0))]
-        [(1, 'POINT (1 1)'), (2, 'POINT (2 2)'), (3, 'POINT (3 3)')]
         """
-        if self._n_geoms == 0:
-            return []
-
         result = []
 
-        def callback(item, userdata):
-            geom = ctypes.cast(item, ctypes.py_object).value
-            result.append(geom)
-
-        lgeos.GEOSSTRtree_query(self._tree_handle, geom._geom, lgeos.GEOSQueryCallback(callback), None)
+        @query_callback
+        def callback(value):
+            result.append(value)
 
+        self.query_cb(geom, callback)
         return result
 
-    def nearest(self, geom):
-        """
-        Get the nearest object in the index to a geometry object.
+    def query_cb(self, geom, callback):
+        """Searches the tree for nodes and applies a callback function.
 
         Parameters
         ----------
-        geom : geometry object
+        geom : Geometry
             The query geometry
+        callback : callable
+            This is a function which takes a node value as argument and
+            which is decorated by shapely.strtree.query_callback. See
+            STRtree.query() for an example.
 
         Returns
         -------
-        geometry object
-            The nearest geometry object in the index to `geom`.
+        None
+
+        """
+        if self._tree is None or not self._initdata:
+            return
+
+        lgeos.GEOSSTRtree_query(
+            self._tree, geom._geom, lgeos.GEOSQueryCallback(callback), None
+        )
+
+    def nearest(self, geom):
+        """Finds the tree node nearest to a given geometry object.
 
-            Will always only return *one* object even if several
-            in the index are the minimum distance away.
+        Parameters
+        ----------
+        geom : geometry object
+            The query geometry
 
-            `None` if the index is empty.
+        Returns
+        -------
+        object or None
+            The value of the tree node nearest to geom or None if the
+            index is empty.
 
         Examples
         --------
@@ -210,27 +265,31 @@ def nearest(self, geom):
 
         Will only return one object:
 
-        >>> tree = STRtree ([Point(0, 0), Point(0, 0)])
+        >>> tree = STRtree([Point(0, 0), Point(0, 0)])
         >>> tree.nearest(Point(0, 0)).wkt
         'POINT (0 0)'
+
         """
-        if self._n_geoms == 0:
+        if self._tree is None or not self._initdata:
             return None
 
+        # In a future version of shapely, geometries will be hashable
+        # and we won't need to reindex like this.
+        geoms = {id(v): g for g, v in self._initdata}
+
+        @nearest_callback
+        def callback(value, geom):
+            value_geom = geoms[id(value)]
+            return geom.distance(value_geom)
+
         envelope = geom.envelope
 
-        def callback(item1, item2, distance, userdata):
-            try:
-                geom1 = ctypes.cast(item1, ctypes.py_object).value
-                geom2 = ctypes.cast(item2, ctypes.py_object).value
-                dist = ctypes.cast(distance, ctypes.POINTER(ctypes.c_double))
-                lgeos.GEOSDistance(geom1._geom, geom2._geom, dist)
-                return 1
-            except Exception:
-                return 0
-
-        item = lgeos.GEOSSTRtree_nearest_generic(self._tree_handle, ctypes.py_object(geom), envelope._geom, \
-            lgeos.GEOSDistanceCallback(callback), None)
-        result = ctypes.cast(item, ctypes.py_object).value
+        item = lgeos.GEOSSTRtree_nearest_generic(
+            self._tree,
+            ctypes.py_object(geom),
+            envelope._geom,
+            lgeos.GEOSDistanceCallback(callback),
+            None,
+        )
 
-        return result
+        return ctypes.cast(item, ctypes.py_object).value