From 582df7ce5e43ef8170999c08d624b79e151b8660 Mon Sep 17 00:00:00 2001
From: Jacopo Farina <jacopo.farina@flixbus.com>
Date: Tue, 14 Sep 2021 11:14:26 +0200
Subject: [PATCH] Add Shapely-based adapter for PostGIS geometry, and docs

Previous commits were messy due to some rebase gone wrong, moved them to
a single commit
---
 .mypy.ini                         |  3 ++
 docs/basic/pgtypes.rst            | 53 +++++++++++++++++++
 psycopg/psycopg/types/geometry.py | 86 +++++++++++++++++++++++++++++++
 3 files changed, 142 insertions(+)
 create mode 100644 psycopg/psycopg/types/geometry.py

diff --git a/.mypy.ini b/.mypy.ini
index 06aa120de..ea6e1e5b4 100644
--- a/.mypy.ini
+++ b/.mypy.ini
@@ -10,3 +10,6 @@ ignore_missing_imports = True
 
 [mypy-setuptools]
 ignore_missing_imports = True
+
+[mypy-shapely.*]
+ignore_missing_imports = True
diff --git a/docs/basic/pgtypes.rst b/docs/basic/pgtypes.rst
index 81f3ffa5d..f8d81c177 100644
--- a/docs/basic/pgtypes.rst
+++ b/docs/basic/pgtypes.rst
@@ -211,3 +211,56 @@ Example::
 
     >>> conn.execute("SELECT 'foo => bar'::hstore").fetchone()[0]
     {'foo': 'bar'}
+
+Geometry adaptation using Shapely
+---------------------------------
+
+When using the PostGIS_ extension, it can be useful to retrieve geometry_
+values and have them automatically converted to Shapely_ instances. Likewise,
+you may want to store such instances in the database and have the conversion
+happen automatically.
+
+To support this, you will need to install Shapely_
+
+.. _PostGIS: https://postgis.net/
+.. _geometry: https://postgis.net/docs/geometry.html
+.. _Shapely: https://github.com/Toblerity/Shapely
+.. _shape: https://shapely.readthedocs.io/en/stable/manual.html#shapely.geometry.shape
+
+Since PostgGIS is an extension, its oid is not well known, so it is necessary
+to use `~psycopg.types.TypeInfo` to query the database and get its oid. After
+that you can use `~psycopg.types.geometry.register_shapely()` to allow dumping
+`shape`_ instances to :sql:`geometry` columns and parsing :sql:`geometry` back
+to `!shape` in the context where it is registered.
+
+.. autofunction:: psycopg.types.geometry.register_shapely
+
+Example::
+
+    >>> from psycopg.types import TypeInfo
+    >>> from psycopg.types.geometry import register_shapely
+    >>> from shapely.geometry import Point
+
+    >>> info = TypeInfo.fetch(conn, "geometry")
+    >>> register_shapely(info, conn)
+
+    >>> conn.execute("SELECT pg_typeof(%s)", [Point(1.2, 3.4)]).fetchone()[0]
+    'geometry'
+
+    >>> conn.execute("""
+    ... SELECT ST_GeomFromGeoJSON('{
+    ...     "type":"Point",
+    ...     "coordinates":[-48.23456,20.12345]}')
+    ... """).fetchone()[0]
+    <shapely.geometry.multipolygon.MultiPolygon object at 0x7fb131f3cd90>
+
+Notice that the adapter is registered on the specific object, other
+connections will be unaffected::
+
+    >>> conn2 = psycopg.connect(CONN_STR)
+    >>> conn2.execute("""
+    ... SELECT ST_GeomFromGeoJSON('{
+    ...     "type":"Point",
+    ...     "coordinates":[-48.23456,20.12345]}')
+    ... """).fetchone()[0]
+    '0101000020E61000009279E40F061E48C0F2B0506B9A1F3440'
diff --git a/psycopg/psycopg/types/geometry.py b/psycopg/psycopg/types/geometry.py
new file mode 100644
index 000000000..7956ebb38
--- /dev/null
+++ b/psycopg/psycopg/types/geometry.py
@@ -0,0 +1,86 @@
+"""
+Adapters for PostGIS geometries
+"""
+
+from typing import Optional, Type
+
+from .. import postgres
+from ..abc import AdaptContext
+from ..adapt import Dumper, Loader
+from ..pq import Format
+from .._typeinfo import TypeInfo
+
+
+try:
+    import shapely.wkb as wkb
+    from shapely.geometry.base import BaseGeometry
+
+except ImportError:
+    raise ImportError(
+        "The module psycopg.types.geometry requires the package 'Shapely'"
+        " to be installed"
+    )
+
+
+class GeometryBinaryLoader(Loader):
+    format = Format.BINARY
+
+    def load(self, data: bytes) -> "BaseGeometry":
+        return wkb.loads(data)
+
+
+class GeometryLoader(Loader):
+    format = Format.TEXT
+
+    def load(self, data: bytes) -> "BaseGeometry":
+        # it's a hex string in binary
+        return wkb.loads(data.decode(), hex=True)
+
+
+class GeometryBinaryDumper(Dumper):
+    format = Format.BINARY
+
+    def dump(self, obj: "BaseGeometry") -> bytes:
+        return wkb.dumps(obj).encode()  # type: ignore
+
+
+class GeometryDumper(Dumper):
+    format = Format.TEXT
+
+    def dump(self, obj: "BaseGeometry") -> bytes:
+        return wkb.dumps(obj, hex=True).encode()  # type: ignore
+
+
+def register_shapely(
+    info: TypeInfo, context: Optional[AdaptContext] = None
+) -> None:
+    """Register Shapely dumper and loaders.
+
+    After invoking this function on an adapter, the queries retrieving
+    PostGIS geometry objects will return Shapely's shape object instances
+    both in text and binary mode.
+
+    Similarly, shape objects can be sent to the database.
+
+    This requires the Shapely library to be installed.
+
+    :param info: The object with the information about the geometry type.
+    :param context: The context where to register the adapters. If `!None`,
+        register it globally.
+
+    """
+
+    info.register(context)
+    adapters = context.adapters if context else postgres.adapters
+    # Generate and register the text and binary dumper
+    binary_dumper: Type[GeometryBinaryDumper] = type(
+        "GeometryBinaryDumper", (GeometryBinaryDumper,), {"oid": info.oid}
+    )
+    dumper: Type[GeometryDumper] = type(
+        "GeometryDumper", (GeometryDumper,), {"oid": info.oid}
+    )
+
+    adapters.register_loader(info.oid, GeometryBinaryLoader)
+    adapters.register_loader(info.oid, GeometryLoader)
+    adapters.register_dumper(BaseGeometry, binary_dumper)
+    adapters.register_dumper(BaseGeometry, dumper)
-- 
2.47.3