From 3a470019a9f5346c7aeeb1640008526c9b821bd8 Mon Sep 17 00:00:00 2001
From: Kamil Monicz <kamil@monicz.dev>
Date: Sun, 23 Mar 2025 10:30:11 +0000
Subject: [PATCH] fix(shapely): add SRID support

---
 psycopg/psycopg/types/shapely.py | 14 +++++---------
 tests/types/test_shapely.py      | 10 +++++++---
 2 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/psycopg/psycopg/types/shapely.py b/psycopg/psycopg/types/shapely.py
index 78260e687..bce1c2490 100644
--- a/psycopg/psycopg/types/shapely.py
+++ b/psycopg/psycopg/types/shapely.py
@@ -12,7 +12,7 @@ from .._compat import cache
 from .._typeinfo import TypeInfo
 
 try:
-    from shapely.wkb import dumps, loads
+    from shapely import from_wkb, to_wkb
     from shapely.geometry.base import BaseGeometry
 
 except ImportError:
@@ -26,29 +26,25 @@ class GeometryBinaryLoader(Loader):
     format = Format.BINARY
 
     def load(self, data: Buffer) -> BaseGeometry:
-        if not isinstance(data, bytes):
-            data = bytes(data)
-        return loads(data)
+        return from_wkb(bytes(data))
 
 
 class GeometryLoader(Loader):
     def load(self, data: Buffer) -> BaseGeometry:
         # it's a hex string in binary
-        if isinstance(data, memoryview):
-            data = bytes(data)
-        return loads(data.decode(), hex=True)
+        return from_wkb(bytes(data))
 
 
 class BaseGeometryBinaryDumper(Dumper):
     format = Format.BINARY
 
     def dump(self, obj: BaseGeometry) -> Buffer | None:
-        return dumps(obj)  # type: ignore
+        return to_wkb(obj, include_srid=True)
 
 
 class BaseGeometryDumper(Dumper):
     def dump(self, obj: BaseGeometry) -> Buffer | None:
-        return dumps(obj, hex=True).encode()  # type: ignore
+        return to_wkb(obj, True, include_srid=True).encode()
 
 
 def register_shapely(info: TypeInfo, context: AdaptContext | None = None) -> None:
diff --git a/tests/types/test_shapely.py b/tests/types/test_shapely.py
index 98fbc2052..ab5d2536e 100644
--- a/tests/types/test_shapely.py
+++ b/tests/types/test_shapely.py
@@ -7,6 +7,7 @@ from psycopg.types import TypeInfo
 
 pytest.importorskip("shapely")
 
+from shapely import get_srid, set_srid
 from shapely.geometry import MultiPolygon, Point, Polygon
 
 from psycopg.types.shapely import register_shapely
@@ -105,7 +106,8 @@ def test_with_adapter(shapely_conn):
 @pytest.mark.parametrize("fmt_out", Format)
 def test_write_read_shape(shapely_conn, fmt_in, fmt_out):
     SAMPLE_POINT = Point(1.2, 3.4)
-    SAMPLE_POLYGON = Polygon([(0, 0), (1, 1), (1, 0)])
+    SAMPLE_POLYGON_4326 = Polygon([(0, 0), (1, 1), (1, 0)])
+    set_srid(SAMPLE_POLYGON_4326, 4326)
 
     with shapely_conn.cursor(binary=fmt_out) as cur:
         cur.execute(
@@ -122,16 +124,18 @@ def test_write_read_shape(shapely_conn, fmt_in, fmt_out):
         )
         cur.execute(
             f"insert into sample_geoms(id, geom) VALUES(2, %{fmt_in})",
-            (SAMPLE_POLYGON,),
+            (SAMPLE_POLYGON_4326,),
         )
 
         cur.execute("select geom from sample_geoms where id=1")
         result = cur.fetchone()[0]
         assert result == SAMPLE_POINT
+        assert get_srid(result) == 0
 
         cur.execute("select geom from sample_geoms where id=2")
         result = cur.fetchone()[0]
-        assert result == SAMPLE_POLYGON
+        assert result == SAMPLE_POLYGON_4326
+        assert get_srid(result) == 4326
 
 
 @pytest.mark.parametrize("fmt_out", Format)
-- 
2.47.2