Added srid keyword to wkb.dumps. Closes #592.

snorfalorpagus · snorfalorpagus · commit 89bd0780b9c7 · 2018-05-28T16:32:01.000+01:00
diff --git a/shapely/wkb.py b/shapely/wkb.py
@@ -1,15 +1,16 @@
 """Load/dump geometries using the well-known binary (WKB) format
 """
 
-from shapely import geos
+from shapely.geos import WKBReader, WKBWriter, lgeos
+from shapely.geometry.base import geom_factory
 
 # Pickle-like convenience functions
 
 def loads(data, hex=False):
     """Load a geometry from a WKB byte string, or hex-encoded string if
     ``hex=True``.
     """
-    reader = geos.WKBReader(geos.lgeos)
+    reader = WKBReader(lgeos)
     if hex:
         return reader.read_hex(data)
     else:
@@ -20,12 +21,29 @@ def load(fp, hex=False):
     data = fp.read()
     return loads(data, hex=hex)
 
-def dumps(ob, hex=False, **kw):
+def dumps(ob, hex=False, srid=None, **kw):
     """Dump a WKB representation of a geometry to a byte string, or a
     hex-encoded string if ``hex=True``.
-
-    See available keyword output settings in ``shapely.geos.WKBWriter``."""
-    writer = geos.WKBWriter(geos.lgeos, **kw)
+    
+    Parameters
+    ----------
+    ob : geometry
+        The geometry to export to well-known binary (WKB) representation
+    hex : bool
+        If true, export the WKB as a hexidecimal string. The default is to
+        return a binary string/bytes object.
+    srid : int
+        Spatial reference system ID to include in the output. The default value
+        means no SRID is included.
+    **kw : kwargs
+        See available keyword output settings in ``shapely.geos.WKBWriter``."""
+    if srid is not None:
+        # clone the object and set the SRID before dumping
+        geom = lgeos.GEOSGeom_clone(ob._geom)
+        lgeos.GEOSSetSRID(geom, srid)
+        ob = geom_factory(geom)
+        kw["include_srid"] = True
+    writer = WKBWriter(lgeos, **kw)
     if hex:
         return writer.write_hex(ob)
     else:
diff --git a/tests/test_wkb.py b/tests/test_wkb.py
@@ -0,0 +1,51 @@
+from shapely.wkb import dumps, loads
+from shapely.geometry import Point
+import binascii
+
+
+def bin2hex(value):
+    return binascii.b2a_hex(value).upper().decode("utf-8")
+
+
+def hex2bin(value):
+    return binascii.a2b_hex(value)
+
+
+def test_dumps_srid():
+    p1 = Point(1.2, 3.4)
+    result = dumps(p1)
+    assert bin2hex(result) == "0101000000333333333333F33F3333333333330B40"
+    result = dumps(p1, srid=4326)
+    assert bin2hex(result) == "0101000020E6100000333333333333F33F3333333333330B40"
+
+
+def test_dumps_endianness():
+    p1 = Point(1.2, 3.4)
+    result = dumps(p1)
+    assert bin2hex(result) == "0101000000333333333333F33F3333333333330B40"
+    result = dumps(p1, big_endian=False)
+    assert bin2hex(result) == "0101000000333333333333F33F3333333333330B40"
+    result = dumps(p1, big_endian=True)
+    assert bin2hex(result) == "00000000013FF3333333333333400B333333333333"
+
+
+def test_dumps_hex():
+    p1 = Point(1.2, 3.4)
+    result = dumps(p1, hex=True)
+    assert result == "0101000000333333333333F33F3333333333330B40"
+    
+
+def test_loads_srid():
+    # load a geometry which includes an srid
+    geom = loads(hex2bin("0101000020E6100000333333333333F33F3333333333330B40"))
+    assert isinstance(geom, Point)
+    assert geom.coords[:] == [(1.2, 3.4)]
+    # by default srid is not exported
+    result = dumps(geom)
+    assert bin2hex(result) == "0101000000333333333333F33F3333333333330B40"
+    # include the srid in the output
+    result = dumps(geom, include_srid=True)
+    assert bin2hex(result) == "0101000020E6100000333333333333F33F3333333333330B40"
+    # replace geometry srid with another
+    result = dumps(geom, srid=27700)
+    assert bin2hex(result) == "0101000020346C0000333333333333F33F3333333333330B40"
