76 lines
2.0 KiB
Python
76 lines
2.0 KiB
Python
|
"""
|
||
|
|
Adapters for PostGIS geometries
|
||
|
|
"""
|
||
|
|
|
||
|
|
from typing import Optional
|
||
|
|
|
||
|
|
from .. import postgres
|
||
|
|
from ..abc import AdaptContext, Buffer
|
||
|
|
from ..adapt import Dumper, Loader
|
||
|
|
from ..pq import Format
|
||
|
|
from .._typeinfo import TypeInfo
|
||
|
|
|
||
|
|
|
||
|
|
try:
|
||
|
|
from shapely.wkb import loads, dumps
|
||
|
|
from shapely.geometry.base import BaseGeometry
|
||
|
|
|
||
|
|
except ImportError:
|
||
|
|
raise ImportError(
|
||
|
|
"The module psycopg.types.shapely requires the package 'Shapely'"
|
||
|
|
" to be installed"
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
class GeometryBinaryLoader(Loader):
|
||
|
|
format = Format.BINARY
|
||
|
|
|
||
|
|
def load(self, data: Buffer) -> "BaseGeometry":
|
||
|
|
if not isinstance(data, bytes):
|
||
|
|
data = bytes(data)
|
||
|
|
return loads(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)
|
||
|
|
|
||
|
|
|
||
|
|
class BaseGeometryBinaryDumper(Dumper):
|
||
|
|
format = Format.BINARY
|
||
|
|
|
||
|
|
def dump(self, obj: "BaseGeometry") -> bytes:
|
||
|
|
return dumps(obj) # type: ignore
|
||
|
|
|
||
|
|
|
||
|
|
class BaseGeometryDumper(Dumper):
|
||
|
|
def dump(self, obj: "BaseGeometry") -> bytes:
|
||
|
|
return dumps(obj, hex=True).encode() # type: ignore
|
||
|
|
|
||
|
|
|
||
|
|
def register_shapely(info: TypeInfo, context: Optional[AdaptContext] = None) -> None:
|
||
|
|
"""Register Shapely dumper and loaders."""
|
||
|
|
|
||
|
|
# A friendly error warning instead of an AttributeError in case fetch()
|
||
|
|
# failed and it wasn't noticed.
|
||
|
|
if not info:
|
||
|
|
raise TypeError("no info passed. Is the 'postgis' extension loaded?")
|
||
|
|
|
||
|
|
info.register(context)
|
||
|
|
adapters = context.adapters if context else postgres.adapters
|
||
|
|
|
||
|
|
class GeometryDumper(BaseGeometryDumper):
|
||
|
|
oid = info.oid
|
||
|
|
|
||
|
|
class GeometryBinaryDumper(BaseGeometryBinaryDumper):
|
||
|
|
oid = info.oid
|
||
|
|
|
||
|
|
adapters.register_loader(info.oid, GeometryBinaryLoader)
|
||
|
|
adapters.register_loader(info.oid, GeometryLoader)
|
||
|
|
# Default binary dump
|
||
|
|
adapters.register_dumper(BaseGeometry, GeometryDumper)
|
||
|
|
adapters.register_dumper(BaseGeometry, GeometryBinaryDumper)
|