From ccba8c7b3937abfcd8b82bf385f66ea08f4ff78b Mon Sep 17 00:00:00 2001 From: Daniele Varrazzo Date: Tue, 25 Mar 2025 19:23:41 +0100 Subject: [PATCH] fix: manage shapely conditional imports explicitly --- psycopg/psycopg/types/shapely.py | 26 ++++++++++++++------------ tests/types/test_shapely.py | 12 +++++------- 2 files changed, 19 insertions(+), 19 deletions(-) diff --git a/psycopg/psycopg/types/shapely.py b/psycopg/psycopg/types/shapely.py index b8e287b11..a6251666f 100644 --- a/psycopg/psycopg/types/shapely.py +++ b/psycopg/psycopg/types/shapely.py @@ -12,46 +12,48 @@ from .._compat import cache from .._typeinfo import TypeInfo try: - from shapely.geometry.base import BaseGeometry + import shapely + except ImportError: raise ImportError( "The module psycopg.types.shapely requires the package 'Shapely'" " to be installed" ) +from shapely.geometry.base import BaseGeometry -try: - from shapely import from_wkb as from_wkb_compat - from shapely import to_wkb as to_wkb_compat -except ImportError: - # Shapely<2 compatibility - from shapely.wkb import dumps as to_wkb_compat # type: ignore[no-redef] - from shapely.wkb import loads as from_wkb_compat # type: ignore[no-redef] +shapely_version = tuple(int(s) for s in shapely.__version__.split(".") if s.isdigit()) + +if shapely_version >= (2, 0): + from shapely import from_wkb, to_wkb +else: + from shapely.wkb import dumps as to_wkb # type: ignore[no-redef] + from shapely.wkb import loads as from_wkb # type: ignore[no-redef] class GeometryBinaryLoader(Loader): format = Format.BINARY def load(self, data: Buffer) -> BaseGeometry: - return from_wkb_compat(bytes(data)) + return from_wkb(bytes(data)) class GeometryLoader(Loader): def load(self, data: Buffer) -> BaseGeometry: # it's a hex string in binary - return from_wkb_compat(bytes(data)) + return from_wkb(bytes(data)) class BaseGeometryBinaryDumper(Dumper): format = Format.BINARY def dump(self, obj: BaseGeometry) -> Buffer | None: - return to_wkb_compat(obj, include_srid=True) + return to_wkb(obj, include_srid=True) class BaseGeometryDumper(Dumper): def dump(self, obj: BaseGeometry) -> Buffer | None: - return to_wkb_compat(obj, True, include_srid=True).encode() + 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 b2c61c036..b4027b0cb 100644 --- a/tests/types/test_shapely.py +++ b/tests/types/test_shapely.py @@ -9,14 +9,12 @@ pytest.importorskip("shapely") from shapely.geometry import MultiPolygon, Point, Polygon -try: - from shapely import get_srid, set_srid -except ImportError: - # Shapely<2 compatibility, no notion of SRID - get_srid = None # type: ignore[assignment] - set_srid = None # type: ignore[assignment] +from psycopg.types.shapely import register_shapely, shapely_version -from psycopg.types.shapely import register_shapely +if shapely_version >= (2, 0): + from shapely import get_srid, set_srid +else: + set_srid = get_srid = None # type: ignore[assignment] pytestmark = [ pytest.mark.postgis, -- 2.47.2