Refactor for asyncpg support

Author: yohanboniface

README.md

@@ -1,25 +1,30 @@
 [![Circle CI](https://img.shields.io/circleci/project/yohanboniface/psycopg-postgis.svg)](https://circleci.com/gh/yohanboniface/psycopg-postgis) [![PyPI](https://img.shields.io/pypi/v/psycopg-postgis.svg)](https://pypi.python.org/pypi/psycopg-postgis) [![PyPI](https://img.shields.io/pypi/pyversions/psycopg-postgis.svg)](https://pypi.python.org/pypi/psycopg-postgis) [![PyPI](https://img.shields.io/pypi/implementation/psycopg-postgis.svg)](https://pypi.python.org/pypi/psycopg-postgis) [![PyPI](https://img.shields.io/pypi/status/psycopg-postgis.svg)](https://pypi.python.org/pypi/psycopg-postgis)
 
-# psycopg-postgis
+# python-postgis
 
-PostGIS helpers for psycopg2.
+PostGIS helpers for psycopg2 and asyncpg.
 
 ## Install
 
-    pip install psycopg-postgis
+    pip install postgis
 
 If you want a compiled version, first install `cython`:
 
     pip install cython
-    pip install psycopg-postgis
+    pip install postgis
 
 
 ## Usage
 
 You need to register the extension:
 
-    > import postgis
-    > postgis.register(mydatabase.get_cursor())
+    # With psycopg2
+    > from postgis.psycopg import register
+    > register(connection)
+
+    # With asyncpg
+    > from postgis.asyncpg import register
+    > await register(connection)
 
 Then you can pass python geometries instance to psycopg:
 
@@ -36,10 +41,10 @@ And retrieve data as python geometries instances:
 ## Example
 
     > import psycopg2
-    > from postgis import register, LineString
+    > from postgis import LineString
+    > from postgis.psycopg import register
     > db = psycopg2.connect(dbname="test")
-    > cursor = db.cursor()
-    > register(cursor)
+    > register(db)
     > cursor.execute('CREATE TABLE IF NOT EXISTS mytable ("geom" geometry(LineString) NOT NULL)')
     > cursor.execute('INSERT INTO mytable (geom) VALUES (%s)', [LineString([(1, 2), (3, 4)], srid=4326)])
     > cursor.execute('SELECT geom FROM mytable LIMIT 1')

circle.yml

@@ -1,6 +1,6 @@
 machine:
     python:
-        version: 3.4.3
+        version: 3.6.1
 
 database:
     pre:
@@ -13,4 +13,4 @@ test:
 
 dependencies:
     pre:
-    -  pip install pytest
+    -  pip install pytest psycopg2 asyncpg cython

postgis/__init__.py

@@ -1,14 +1,22 @@
-"Postgis helpers for psycopg2."
-from .geometry import Geometry, register
-from .point import Point
+"Postgis helpers for psycopg2 and asyncpg."
+from .geometry import Geometry
+from .geometrycollection import GeometryCollection
 from .linestring import LineString
-from .polygon import Polygon
-from .multipoint import MultiPoint
 from .multilinestring import MultiLineString
+from .multipoint import MultiPoint
 from .multipolygon import MultiPolygon
-from .geometrycollection import GeometryCollection
-from .__meta__ import __version__
+from .point import Point
+from .polygon import Polygon
+from .psycopg import register  # Retrocompat.
 
 __all__ = ['Geometry', 'register', 'Point', 'LineString', 'Polygon',
            'MultiPoint', 'MultiLineString', 'MultiPolygon',
-           'GeometryCollection', '__version__']
+           'GeometryCollection']
+
+try:
+    import pkg_resources
+except ImportError:  # pragma: no cover
+    pass
+else:
+    if __package__:
+        VERSION = pkg_resources.get_distribution(__package__).version

postgis/__meta__.py

@@ -0,0 +1,17 @@
+from .geometry import Geometry
+
+
+async def register(connection):
+
+    def encoder(value):
+        if not isinstance(value, Geometry):
+            raise ValueError('Geometry value must subclass Geometry class')
+        return value.to_ewkb()
+
+    def decoder(value):
+        return Geometry.from_ewkb(value)
+
+    await connection.set_type_codec(
+        'geography', encoder=encoder, decoder=decoder)
+    await connection.set_type_codec(
+        'geometry', encoder=encoder, decoder=decoder)

postgis/asyncpg.py

@@ -0,0 +1,101 @@
+import binascii
+from io import BytesIO
+import struct
+
+
+class Typed(type):
+
+    types = {}
+
+    def __new__(mcs, name, bases, attrs, **kwargs):
+        cls = super().__new__(mcs, name, bases, attrs, **kwargs)
+        if hasattr(cls, 'TYPE'):
+            Typed.types[cls.TYPE] = cls
+        return cls
+
+    def __call__(cls, *args, **kwargs):
+        # Allow to pass an instance as first argument, for blind casting.
+        if args and isinstance(args[0], cls):
+            return args[0]
+        return super().__call__(*args, **kwargs)
+
+
+class Reader:
+
+    __slots__ = ['stream', 'endianness', 'has_z', 'has_m']
+
+    def __init__(self, stream):
+        self.stream = stream
+
+    def clone(self):
+        return type(self)(self.stream)
+
+    def read(self):
+        # https://en.wikipedia.org/wiki/Well-known_text#Well-known_binary
+        byte_order = self.stream.read(1)
+        if byte_order == b'\x00':
+            self.endianness = b'>'
+        elif byte_order == b'\x01':
+            self.endianness = b'<'
+        else:
+            raise Exception('invalid encoding')
+
+        type_ = self.read_int()
+        self.has_z = bool(type_ & 0x80000000)
+        self.has_m = bool(type_ & 0x40000000)
+        srid = self.read_int() if bool(type_ & 0x20000000) else None
+        type_ &= 0x1fffffff
+
+        try:
+            class_ = Typed.types[type_]
+        except KeyError:
+            raise ValueError('unsupported geometry type {0}'.format(type_))
+        else:
+            return class_.from_ewkb_body(self, srid)
+
+    def read_int(self):
+        return struct.unpack(self.endianness + b'I', self.stream.read(4))[0]
+
+    def read_double(self):
+        return struct.unpack(self.endianness + b'd', self.stream.read(8))[0]
+
+    @classmethod
+    def from_hex(cls, value):
+        return cls(BytesIO(binascii.a2b_hex(value))).read()
+
+
+class Writer:
+
+    __slots__ = ['stream']
+
+    def __init__(self, geometry, stream=None):
+        self.stream = stream or BytesIO()
+        try:
+            type_ = geometry.TYPE
+        except AttributeError:
+            raise ValueError('Unknown geometry {}'.format(geometry.__class__))
+
+        # Little endian.
+        self.stream.write(b'\x01')
+        self.write_int(
+            type_ |
+            (0x80000000 if geometry.has_z else 0) |
+            (0x40000000 if geometry.has_m else 0) |
+            (0x20000000 if geometry.has_srid else 0))
+        if geometry.has_srid:
+            self.write_int(geometry.srid)
+
+    def write_int(self, value):
+        self.stream.write(struct.pack(b'<I', value))
+
+    def write_double(self, value):
+        self.stream.write(struct.pack(b'<d', value))
+
+    def clone(self, geometry):
+        return type(self)(geometry, self.stream)
+
+    @classmethod
+    def to_hex(cls, value):
+        writer = cls(value)
+        value.write_ewkb_body(writer)
+        return binascii.b2a_hex(writer.stream.getvalue()).upper()

postgis/ewkb.py

@@ -1,5 +1,11 @@
-from psycopg2 import extensions as _ext
-from .reader import EWKBReader, Typed
+try:
+    # Do not make psycopg2 a requirement.
+    from psycopg2.extensions import ISQLQuote
+except ImportError:
+    print('psycopg2 not installed')
+
+
+from .ewkb import Reader, Typed, Writer
 from .geojson import GeoJSON
 
 
@@ -21,14 +27,25 @@ def has_srid(self):
     def from_ewkb(value, cursor=None):
         if not value:
             return None
-        return EWKBReader.from_hex(value)
+        return Reader.from_hex(value)
+
+    def to_ewkb(self):
+        return Writer.to_hex(self).decode()
+
+    def write_ewkb(self, writer):
+        self.write_ewkb_body(writer.clone(self))
+
+    def text(self):
+        return "ST_GeometryFromText('{}', {})".format(self.wkt, self.srid)
 
+    # Psycopg2 interface.
     def __conform__(self, protocol):
-        if protocol is _ext.ISQLQuote:
+        if protocol is ISQLQuote:
             return self
 
     def getquoted(self):
-        return "ST_GeometryFromText('{}', {})".format(self.wkt, self.srid)
+        return "'{}'".format(self.to_ewkb())
+    # End Psycopg2 interface.
 
     def __str__(self):
         return self.wkt
@@ -55,15 +72,3 @@ def geojson(self):
             'type': self.name,
             'coordinates': self.coords
         })
-
-
-def register(cursor):
-    cursor.execute("SELECT NULL::geometry")
-    oid = cursor.description[0][1]
-    GEOMETRY = _ext.new_type((oid, ), "GEOMETRY", Geometry.from_ewkb)
-    _ext.register_type(GEOMETRY)
-
-    cursor.execute("SELECT NULL::geography")
-    oid = cursor.description[0][1]
-    GEOGRAPHY = _ext.new_type((oid, ), "GEOGRAPHY", Geometry.from_ewkb)
-    _ext.register_type(GEOGRAPHY)

postgis/geometry.py

@@ -37,6 +37,11 @@ def __getitem__(self, item):
     def from_ewkb_body(cls, reader, srid=None):
         return cls([reader.read() for index in range(reader.read_int())], srid)
 
+    def write_ewkb_body(self, writer):
+        writer.write_int(len(self.geoms))
+        for geom in self:
+            geom.write_ewkb(writer)
+
     @property
     def wkt_coords(self):
         return ', '.join(g.wkt for g in self)

postgis/geometrycollection.py

@@ -1,42 +1,18 @@
-from .geometry import Geometry
 from .point import Point
+from .multi import Multi
 
 
-class LineString(Geometry):
+class LineString(Multi):
 
     TYPE = 2
-
-    def __init__(self, points, srid=None):
-        self.points = list(points)
-        if srid:
-            self.srid = srid
-
-    def __iter__(self):
-        for point in self.points:
-            if not isinstance(point, Point):
-                point = Point(*point)
-            yield point
-
-    @property
-    def has_z(self):
-        return self.points[0].has_z
-
-    @property
-    def has_m(self):
-        return self.points[0].has_m
-
-    def __getitem__(self, item):
-        return self.points[item]
+    SUBCLASS = Point
 
     @classmethod
     def from_ewkb_body(cls, reader, srid=None):
         return cls([Point.from_ewkb_body(reader)
                    for index in range(reader.read_int())], srid)
 
-    @property
-    def wkt_coords(self):
-        return ', '.join(p.wkt_coords for p in self)
-
-    @property
-    def coords(self):
-        return tuple(p.coords for p in self)
+    def write_ewkb_body(self, writer):
+        writer.write_int(len(self.geoms))
+        for geom in self:
+            geom.write_ewkb_body(writer)

postgis/linestring.py

@@ -3,18 +3,16 @@
 
 class Multi(Geometry):
 
+    __slots__ = ['geoms', 'srid']
     SUBCLASS = None
 
     def __init__(self, geoms, srid=None):
-        self.geoms = list(geoms)
+        self.geoms = [self.SUBCLASS(g, srid=srid) for g in geoms]
         if srid:
             self.srid = srid
 
     def __iter__(self):
-        for geom in self.geoms:
-            if not isinstance(geom, self.SUBCLASS):
-                geom = self.SUBCLASS(geom)
-            yield geom
+        return iter(self.geoms)
 
     @property
     def has_z(self):
@@ -35,6 +33,11 @@ def from_ewkb_body(cls, reader, srid=None):
     def wkt_coords(self):
         return ', '.join('({})'.format(g.wkt_coords) for g in self)
 
+    def write_ewkb_body(self, writer):
+        writer.write_int(len(self.geoms))
+        for geom in self:
+            geom.write_ewkb(writer)
+
     @property
     def coords(self):
         return tuple(g.coords for g in self)