Add compat for unpickling shapely<2 geometries

CHANGES.txt

@@ -1,6 +1,13 @@
 Changes
 =======
 
+
+2.0rc4 (2022-10-26)
+-------------------
+
+- Added temporary for unpickling shapely<2.0 geometries.
+
+
 2.0rc1 (2022-10-26)
 -------------------
 

src/pygeom.c

@@ -38,14 +38,14 @@ PyObject* GeometryObject_FromGEOS(GEOSGeometry* ptr, GEOSContextHandle_t ctx) {
   } else {
     self->ptr = ptr;
     self->ptr_prepared = NULL;
-    self->weakreflist = (PyObject *)NULL;
+    self->weakreflist = (PyObject*)NULL;
     return (PyObject*)self;
   }
 }
 
 static void GeometryObject_dealloc(GeometryObject* self) {
   if (self->weakreflist != NULL) {
-    PyObject_ClearWeakRefs((PyObject *)self);
+    PyObject_ClearWeakRefs((PyObject*)self);
   }
   if (self->ptr != NULL) {
     // not using GEOS_INIT, but using global context instead
@@ -291,7 +291,87 @@ static PyObject* GeometryObject_richcompare(GeometryObject* self, PyObject* othe
   return result;
 }
 
+
+static PyObject* GeometryObject_SetState(PyObject* self, PyObject* value) {
+  unsigned char* wkb = NULL;
+  Py_ssize_t size;
+  GEOSGeometry* geom = NULL;
+  GEOSWKBReader* reader = NULL;
+
+  PyErr_WarnFormat(PyExc_UserWarning, 0,
+                   "Unpickling a shapely <2.0 geometry object. Please save the pickle "
+                   "again; shapely 2.1 will not have this compatibility.");
+
+  /* Cast the PyObject bytes to char */
+  if (!PyBytes_Check(value)) {
+    PyErr_Format(PyExc_TypeError, "Expected bytes, found %s", value->ob_type->tp_name);
+    return NULL;
+  }
+  size = PyBytes_Size(value);
+  wkb = (unsigned char*)PyBytes_AsString(value);
+  if (wkb == NULL) {
+    return NULL;
+  }
+
+  PyObject* linearring_type_obj = PyList_GET_ITEM(geom_registry[0], 2);
+  if (linearring_type_obj == NULL) {
+    return NULL;
+  }
+  if (!PyType_Check(linearring_type_obj)) {
+    PyErr_Format(PyExc_RuntimeError, "Invalid registry value");
+    return NULL;
+  }
+  PyTypeObject* linearring_type = (PyTypeObject*)linearring_type_obj;
+
+  GEOS_INIT;
+
+  reader = GEOSWKBReader_create_r(ctx);
+  if (reader == NULL) {
+    errstate = PGERR_GEOS_EXCEPTION;
+    goto finish;
+  }
+  geom = GEOSWKBReader_read_r(ctx, reader, wkb, size);
+  if (geom == NULL) {
+    errstate = PGERR_GEOS_EXCEPTION;
+    goto finish;
+  }
+  if (Py_TYPE(self) == linearring_type) {
+    const GEOSCoordSequence* coord_seq = GEOSGeom_getCoordSeq_r(ctx, geom);
+    if (coord_seq == NULL) {
+      errstate = PGERR_GEOS_EXCEPTION;
+      goto finish;
+    }
+    geom = GEOSGeom_createLinearRing_r(ctx, (GEOSCoordSequence*)coord_seq);
+    if (geom == NULL) {
+      errstate = PGERR_GEOS_EXCEPTION;
+      goto finish;
+    }
+  }
+
+  if (((GeometryObject*)self)->ptr != NULL) {
+    GEOSGeom_destroy_r(ctx, ((GeometryObject*)self)->ptr);
+  }
+  ((GeometryObject*)self)->ptr = geom; 
+
+finish:
+
+  if (reader != NULL) {
+    GEOSWKBReader_destroy_r(ctx, reader);
+  }
+
+  GEOS_FINISH;
+
+  if (errstate == PGERR_SUCCESS) {
+    Py_INCREF(Py_None);
+    return Py_None;
+  }
+  return NULL;
+}
+
+
 static PyMethodDef GeometryObject_methods[] = {
+    {"__setstate__", (PyCFunction)GeometryObject_SetState, METH_O,
+     "For unpickling pre-shapely 2.0 pickles"},
     {NULL} /* Sentinel */
 };
 

tests/test_pickle.py

@@ -1,24 +1,71 @@
+import pathlib
+import pickle
+from pickle import dumps, loads, HIGHEST_PROTOCOL
+import warnings
+
+import shapely
+from shapely.geometry import Point, LineString, LinearRing, Polygon, MultiLineString, MultiPoint, MultiPolygon, GeometryCollection, box
+from shapely import wkt
+
 import pytest
-from shapely.geometry import Point, LineString, LinearRing, Polygon, MultiPoint
 
-from pickle import dumps, loads, HIGHEST_PROTOCOL
+
+HERE = pathlib.Path(__file__).parent
+
 
 TEST_DATA = {
-    "point2d": (Point, [(1.0, 2.0)]),
-    "point3d": (Point, [(1.0, 2.0, 3.0)]),
-    "linestring": (LineString, [(0.0, 0.0), (0.0, 1.0), (1.0, 1.0)]),
-    "linearring": (LinearRing, [(0.0, 0.0), (1.0, 0.0), (1.0, 1.0), (0.0, 0.0)]),
-    "polygon": (Polygon, [(0.0, 0.0), (1.0, 0.0), (1.0, 1.0), (0.0, 0.0)]),
-    "multipoint": (MultiPoint, [(1.0, 2.0), (3.0, 4.0), (5.0, 6.0)]),
+    "point2d": Point([(1.0, 2.0)]),
+    "point3d": Point([(1.0, 2.0, 3.0)]),
+    "linestring": LineString([(0.0, 0.0), (0.0, 1.0), (1.0, 1.0)]),
+    "linearring": LinearRing([(0.0, 0.0), (1.0, 0.0), (1.0, 1.0), (0.0, 0.0)]),
+    "polygon": Polygon([(0.0, 0.0), (1.0, 0.0), (1.0, 1.0), (0.0, 0.0)]),
+    "multipoint": MultiPoint([(1.0, 2.0), (3.0, 4.0), (5.0, 6.0)]),
+    "multilinestring": MultiLineString([[(0.0, 0.0), (1.0, 1.0)], [(1.0, 2.0), (3.0, 3.0)]]),
+    "multipolygon": MultiPolygon([box(0, 0, 1, 1), box(2, 2, 3, 3)]),
+    "geometrycollection": GeometryCollection([Point(1.0, 2.0), box(0, 0, 1, 1)]),
+    "emptypoint": wkt.loads("POINT EMPTY"),
+    "emptypolygon": wkt.loads("POLYGON EMPTY"),
 }
-TEST_NAMES, TEST_DATA = zip(*TEST_DATA.items())
-@pytest.mark.parametrize("cls,coords", TEST_DATA, ids=TEST_NAMES)
-def test_pickle_round_trip(cls, coords):
-    geom1 = cls(coords)
-    assert geom1.has_z == (len(coords[0]) == 3)
+TEST_NAMES, TEST_GEOMS = zip(*TEST_DATA.items())
+
+
+@pytest.mark.parametrize("geom1", TEST_GEOMS, ids=TEST_NAMES)
+def test_pickle_round_trip(geom1):
     data = dumps(geom1, HIGHEST_PROTOCOL)
-    geom2 = loads(data)
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        geom2 = loads(data)
     assert geom2.has_z == geom1.has_z
     assert type(geom2) is type(geom1)
     assert geom2.geom_type == geom1.geom_type
     assert geom2.wkt == geom1.wkt
+
+
+@pytest.mark.parametrize("fname", (HERE / "data").glob("*.pickle"), ids=lambda fname: fname.name)
+def test_unpickle_pre_20(fname):
+    from shapely.testing import assert_geometries_equal
+
+    geom_type = fname.name.split("_")[0]
+    expected = TEST_DATA[geom_type]
+
+    with open(fname, "rb") as f:
+        with pytest.warns(UserWarning):
+            result = pickle.load(f)
+
+    assert_geometries_equal(result, expected)
+
+
+if __name__ == "__main__":
+    datadir = HERE / "data"
+    datadir.mkdir(exist_ok=True)
+
+    shapely_version = shapely.__version__
+    print(shapely_version)
+    print(shapely.geos.geos_version)
+
+    for name, geom in TEST_DATA.items():
+        if name == "emptypoint" and shapely.geos.geos_version < (3, 9, 0):
+            # Empty Points cannot be represented in WKB
+            continue
+        with open(datadir / f"{name}_{shapely_version}.pickle", "wb") as f:
+            pickle.dump(geom, f)