ENH: Implement from_shapely (#61)

pygeos/__init__.py

@@ -1,6 +1,7 @@
 from .lib import GEOSException  # NOQA
 from .lib import Geometry  # NOQA
-from .lib import geos_version  # NOQA
+from .lib import geos_version, geos_version_string  # NOQA
+from .lib import geos_capi_version, geos_capi_version_string  # NOQA
 from .geometry import *
 from .creation import *
 from .constructive import *

pygeos/io.py

@@ -2,9 +2,17 @@
 
 from . import Geometry  # noqa
 from . import lib
+from . import geos_capi_version_string
 
+try:
+    from shapely.geos import geos_version_string as shapely_geos_version
+    from shapely.geometry.base import BaseGeometry as ShapelyGeometry
+except ImportError:
+    shapely_geos_version = None
+    ShapelyGeometry = None
 
-__all__ = ["from_wkb", "from_wkt", "to_wkb", "to_wkt"]
+
+__all__ = ["from_shapely", "from_wkb", "from_wkt", "to_wkb", "to_wkt"]
 
 
 def to_wkt(
@@ -80,7 +88,9 @@ def to_wkt(
     )
 
 
-def to_wkb(geometry, hex=False, output_dimension=3, byte_order=-1, include_srid=False, **kwargs):
+def to_wkb(
+    geometry, hex=False, output_dimension=3, byte_order=-1, include_srid=False, **kwargs
+):
     r"""
     Converts to the Well-Known Binary (WKB) representation of a Geometry.
 
@@ -172,3 +182,42 @@ def from_wkb(geometry, **kwargs):
     # of array elements)
     geometry = np.asarray(geometry, dtype=object)
     return lib.from_wkb(geometry, **kwargs)
+
+
+def from_shapely(geometry, **kwargs):
+    """Creates geometries from shapely Geometry objects.
+
+    This function requires the GEOS version of PyGEOS and shapely to be equal.
+
+    Parameters
+    ----------
+    geometry : shapely Geometry object or array_like
+
+    Examples
+    --------
+    >>> from shapely.geometry import Point   # doctest: +SKIP
+    >>> from_shapely(Point(1, 2))   # doctest: +SKIP
+    <pygeos.Geometry POINT (1 2)>
+    """
+    if shapely_geos_version is None:
+        raise ImportError("This function requires shapely")
+
+    # shapely has something like: "3.6.2-CAPI-1.10.2 4d2925d6"
+    # pygeos has something like: "3.6.2-CAPI-1.10.2"
+    if not shapely_geos_version.startswith(geos_capi_version_string):
+        raise ImportError(
+            "The shapely GEOS version ({}) is incompatible with the GEOS "
+            "version PyGEOS was compiled with ({})".format(
+                shapely_geos_version, geos_capi_version_string
+            )
+        )
+
+    if isinstance(geometry, ShapelyGeometry):
+        # this so that the __array_interface__ of the shapely geometry is not
+        # used, converting the Geometry to its coordinates
+        arr = np.empty(1, dtype=object)
+        arr[0] = geometry
+        arr.shape = ()
+    else:
+        arr = np.asarray(geometry, dtype=object)
+    return lib.from_shapely(arr, **kwargs)

pygeos/test/test_io.py

@@ -1,6 +1,7 @@
 import numpy as np
 import pygeos
 import pytest
+from unittest import mock
 
 from .common import all_types, point
 
@@ -10,6 +11,16 @@
 )
 
 
+class ShapelyGeometryMock:
+    def __init__(self, g):
+        self.g = g
+        self.__geom__ = g._ptr if hasattr(g, "_ptr") else g
+
+    def __array_interface__(self):
+        # this should not be called
+        raise NotImplementedError()
+
+
 def test_from_wkt():
     expected = pygeos.points(1, 1)
     actual = pygeos.from_wkt("POINT (1 1)")
@@ -165,7 +176,9 @@ def test_to_wkb_hex():
 def test_to_wkb_3D():
     point_z = pygeos.points(1, 1, 1)
     actual = pygeos.to_wkb(point_z)
+    # fmt: off
     assert actual == b"\x01\x01\x00\x00\x80\x00\x00\x00\x00\x00\x00\xf0?\x00\x00\x00\x00\x00\x00\xf0?\x00\x00\x00\x00\x00\x00\xf0?"  # noqa
+    # fmt: on
     actual = pygeos.to_wkb(point_z, output_dimension=2)
     assert actual == POINT11_WKB
 
@@ -209,3 +222,42 @@ def test_to_wkb_srid():
     point_with_srid = pygeos.set_srid(point, np.int32(4326))
     result = pygeos.to_wkb(point_with_srid, include_srid=True)
     assert np.frombuffer(result[5:9], "<u4").item() == 4326
+
+
+@pytest.mark.parametrize("geom", all_types)
+@mock.patch("pygeos.io.ShapelyGeometry", ShapelyGeometryMock)
+@mock.patch("pygeos.io.shapely_geos_version", pygeos.geos_capi_version_string)
+def test_from_shapely(geom):
+    actual = pygeos.from_shapely(ShapelyGeometryMock(geom))
+    assert isinstance(geom, pygeos.Geometry)
+    assert pygeos.equals(geom, actual)
+    assert geom._ptr != actual._ptr
+
+
+@mock.patch("pygeos.io.ShapelyGeometry", ShapelyGeometryMock)
+@mock.patch("pygeos.io.shapely_geos_version", pygeos.geos_capi_version_string)
+def test_from_shapely_arr():
+    actual = pygeos.from_shapely([ShapelyGeometryMock(point), None])
+    assert pygeos.equals(point, actual[0])
+
+
+@mock.patch("pygeos.io.ShapelyGeometry", ShapelyGeometryMock)
+@mock.patch("pygeos.io.shapely_geos_version", pygeos.geos_capi_version_string)
+def test_from_shapely_none():
+    actual = pygeos.from_shapely(None)
+    assert actual is None
+
+
+@pytest.mark.parametrize("geom", [1, 2.3, "x", ShapelyGeometryMock(None)])
+@mock.patch("pygeos.io.ShapelyGeometry", ShapelyGeometryMock)
+@mock.patch("pygeos.io.shapely_geos_version", pygeos.geos_capi_version_string)
+def test_from_shapely_error(geom):
+    with pytest.raises(TypeError):
+        pygeos.from_shapely(geom)
+
+
+# We have >= 3.5 in PyGEOS. Test with some random older version.
+@mock.patch("pygeos.io.shapely_geos_version", "2.3.4-abc")
+def test_from_shapely_incompatible_versions():
+    with pytest.raises(ImportError):
+        pygeos.from_shapely(point)

pygeos/test/test_misc.py

@@ -6,4 +6,12 @@ def test_version():
 
 
 def test_geos_version():
-    assert isinstance(pygeos.geos_version, str)
+    expected = "{0:d}.{1:d}.{2:d}".format(*pygeos.geos_version)
+    assert pygeos.geos_version_string == expected
+
+
+def test_geos_capi_version():
+    expected = "{0:d}.{1:d}.{2:d}-CAPI-{3:d}.{4:d}.{5:d}".format(
+        *(pygeos.geos_version + pygeos.geos_capi_version)
+    )
+    assert pygeos.geos_capi_version_string == expected

src/lib.c

@@ -62,8 +62,20 @@ PyMODINIT_FUNC PyInit_lib(void)
     import_array();
     import_umath();
 
-    /* export the version as a python string */
-    PyModule_AddObject(m, "geos_version", PyUnicode_FromString(GEOS_VERSION));
+    /* export the GEOS versions as python tuple and string */
+    PyModule_AddObject(m, "geos_version", PyTuple_Pack(3,
+        PyLong_FromLong((long) GEOS_VERSION_MAJOR),
+        PyLong_FromLong((long) GEOS_VERSION_MINOR),
+        PyLong_FromLong((long) GEOS_VERSION_PATCH)
+    ));
+    PyModule_AddObject(m, "geos_capi_version", PyTuple_Pack(3,
+        PyLong_FromLong((long) GEOS_CAPI_VERSION_MAJOR),
+        PyLong_FromLong((long) GEOS_CAPI_VERSION_MINOR),
+        PyLong_FromLong((long) GEOS_CAPI_VERSION_PATCH)
+    ));
+
+    PyModule_AddObject(m, "geos_version_string", PyUnicode_FromString(GEOS_VERSION));
+    PyModule_AddObject(m, "geos_capi_version_string", PyUnicode_FromString(GEOS_CAPI_VERSION));
 
     if (init_ufuncs(m, d) < 0) {
         return NULL;

src/pygeom.c

@@ -84,12 +84,8 @@ static PyObject *GeometryObject_FromWKT(PyTypeObject *type, PyObject *value)
     GEOSGeometry *geom;
     GEOSWKTReader *reader;
 
-    /* Cast the PyObject (bytes or str) to char* */
-    if (PyBytes_Check(value)) {
-        wkt = PyBytes_AsString(value);
-        if (wkt == NULL) { return NULL; }
-    }
-    else if (PyUnicode_Check(value)) {
+    /* Cast the PyObject str to char* */
+    if (PyUnicode_Check(value)) {
         wkt = PyUnicode_AsUTF8(value);
         if (wkt == NULL) { return NULL; }
     } else {

src/ufuncs.c

@@ -1127,6 +1127,47 @@ static void from_wkt_func(char **args, npy_intp *dimensions,
 }
 static PyUFuncGenericFunction from_wkt_funcs[1] = {&from_wkt_func};
 
+static char from_shapely_dtypes[2] = {NPY_OBJECT, NPY_OBJECT};
+static void from_shapely_func(char **args, npy_intp *dimensions,
+                              npy_intp *steps, void *data)
+{
+    void *context_handle = geos_context[0];
+    GEOSGeometry *in_ptr, *ret_ptr;
+    PyObject *in1, *attr;
+
+    UNARY_LOOP {
+        /* ip1 is pointer to array element PyObject* */
+        in1 = *(PyObject **)ip1;
+
+        if (in1 == Py_None) {
+            /* None in the input propagates to the output */
+            ret_ptr = NULL;
+        }
+        else {
+            /* Get the __geom__ attribute */
+            attr = PyObject_GetAttrString(in1, "__geom__");
+            if (attr == NULL) {
+                /* Raise if __geom__ does not exist */
+                PyErr_Format(PyExc_TypeError, "Expected a shapely object or None, got %s", Py_TYPE(in1)->tp_name);
+                return;
+            } else if (!PyLong_Check(attr)) {
+                /* Raise if __geom__ is of incorrect type */
+                PyErr_Format(PyExc_TypeError, "Expected int for the __geom__ attribute, got %s", Py_TYPE(attr)->tp_name);
+                Py_XDECREF(attr);
+                return;
+            }
+            /* Convert it to a GEOSGeometry pointer */
+            in_ptr = PyLong_AsVoidPtr(attr);
+            Py_XDECREF(attr);
+            /* Clone the geometry and finish */
+            ret_ptr = GEOSGeom_clone_r(context_handle, in_ptr);
+            if (ret_ptr == NULL) { return; }
+        }
+        OUTPUT_Y;
+    }
+}
+static PyUFuncGenericFunction from_shapely_funcs[1] = {&from_shapely_func};
+
 static char to_wkb_dtypes[6] = {NPY_OBJECT, NPY_BOOL, NPY_INT, NPY_INT, NPY_BOOL, NPY_OBJECT};
 static void to_wkb_func(char **args, npy_intp *dimensions,
                         npy_intp *steps, void *data)
@@ -1397,6 +1438,7 @@ int init_ufuncs(PyObject *m, PyObject *d)
     DEFINE_CUSTOM (from_wkt, 1);
     DEFINE_CUSTOM (to_wkb, 5);
     DEFINE_CUSTOM (to_wkt, 5);
+    DEFINE_CUSTOM (from_shapely, 1);
 
     Py_DECREF(ufunc);
     return 0;