Fix segfault when using unpickled STRtree in another process (#915)

* Add test for STRtree unpickling segfault

* Fix segfault when using unpickled STRtree in another process

shapely/strtree.py

@@ -36,13 +36,26 @@ def __init__(self, geoms):
         # 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):
         # GEOS STRtree capacity has to be > 1
         self._tree_handle = lgeos.GEOSSTRtree_create(max(2, len(geoms)))
         for geom in geoms:
             lgeos.GEOSSTRtree_insert(self._tree_handle, geom._geom, ctypes.py_object(geom))
 
-        # Keep references to geoms.
-        self._geoms = list(geoms)
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        del state["_tree_handle"]
+        return state
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+        self._init_tree_handle(self._geoms)
 
     def __del__(self):
         if self._tree_handle is not None:

tests/test_strtree.py

@@ -1,4 +1,8 @@
 import gc
+import pickle
+import subprocess
+import sys
+from pathlib import Path
 
 from shapely.strtree import STRtree
 from shapely.geometry import Point, Polygon
@@ -79,3 +83,19 @@ def test_safe_delete():
     del tree
 
     strtree.lgeos = _lgeos
+
+
+@requires_geos_342
+def test_pickle_persistence():
+    """
+    Don't crash trying to use unpickled GEOS handle.
+    """
+    tree = STRtree([Point(i, i).buffer(0.1) for i in range(3)])
+    pickled_strtree = pickle.dumps(tree)
+    print("pickled strtree:", repr(pickled_strtree))
+    unpickle_script_file_path = Path(__file__).parent / "unpickle-strtree.py"
+    subprocess.run(
+        [sys.executable, str(unpickle_script_file_path)],
+        input=pickled_strtree,
+        check=True
+    )

tests/unpickle-strtree.py

@@ -0,0 +1,24 @@
+# See test_strtree.py::test_pickle_persistence
+
+import sys
+import os
+sys.path.append(os.getcwd())
+
+import pickle
+from shapely.geometry import Point
+from shapely.geos import geos_version
+
+
+if __name__ == "__main__":
+    pickled_strtree = sys.stdin.buffer.read()
+    print("received pickled strtree:", repr(pickled_strtree))
+    strtree = pickle.loads(pickled_strtree)
+    # Exercise API.
+    print("calling \"query()\"...")
+    strtree.query(Point(0, 0))
+    if geos_version >= (3, 6, 0):
+        print("calling \"nearest()\"...")
+        strtree.nearest(Point(0, 0))
+    else:
+        print("skipping \"nearest()\"")
+    print("done")