Fixed #26510 -- Allowed dim/trim/precision as WKTWriter init arguments

Thanks Tim Graham for the review.

django/contrib/gis/geos/geometry.py

@@ -408,7 +408,7 @@ class GEOSGeometry(GEOSBase, ListMixin):
     @property
     def wkt(self):
         "Returns the WKT (Well-Known Text) representation of this Geometry."
-        return wkt_w(3 if self.hasz else 2).write(self).decode()
+        return wkt_w(dim=3 if self.hasz else 2).write(self).decode()
 
     @property
     def hex(self):
@@ -419,7 +419,7 @@ class GEOSGeometry(GEOSBase, ListMixin):
         """
         # A possible faster, all-python, implementation:
         #  str(self.wkb).encode('hex')
-        return wkb_w(3 if self.hasz else 2).write_hex(self)
+        return wkb_w(dim=3 if self.hasz else 2).write_hex(self)
 
     @property
     def hexewkb(self):
@@ -428,7 +428,7 @@ class GEOSGeometry(GEOSBase, ListMixin):
         extension of the WKB specification that includes SRID value that are
         a part of this geometry.
         """
-        return ewkb_w(3 if self.hasz else 2).write_hex(self)
+        return ewkb_w(dim=3 if self.hasz else 2).write_hex(self)
 
     @property
     def json(self):

django/contrib/gis/geos/prototypes/io.py

@@ -170,6 +170,14 @@ class WKTWriter(IOBase):
     _trim = False
     _precision = None
 
+    def __init__(self, dim=2, trim=False, precision=None):
+        super(WKTWriter, self).__init__()
+        if bool(trim) != self._trim:
+            self.trim = trim
+        if precision is not None:
+            self.precision = precision
+        self.outdim = dim
+
     def write(self, geom):
         "Returns the WKT representation of the given geometry."
         return wkt_writer_write(self.ptr, geom.ptr)
@@ -212,6 +220,10 @@ class WKBWriter(IOBase):
     _destructor = wkb_writer_destroy
     ptr_type = WKB_WRITE_PTR
 
+    def __init__(self, dim=2):
+        super(WKBWriter, self).__init__()
+        self.outdim = dim
+
     def write(self, geom):
         "Returns the WKB representation of the given geometry."
         return six.memoryview(wkb_writer_write(self.ptr, geom.ptr, byref(c_size_t())))
@@ -280,10 +292,13 @@ def wkt_r():
     return thread_context.wkt_r
 
 
-def wkt_w(dim=2):
+def wkt_w(dim=2, trim=False, precision=None):
     if not thread_context.wkt_w:
-        thread_context.wkt_w = WKTWriter()
+        thread_context.wkt_w = WKTWriter(dim=dim, trim=trim, precision=precision)
+    else:
         thread_context.wkt_w.outdim = dim
+        thread_context.wkt_w.trim = trim
+        thread_context.wkt_w.precision = precision
     return thread_context.wkt_w
 
 
@@ -295,14 +310,16 @@ def wkb_r():
 
 def wkb_w(dim=2):
     if not thread_context.wkb_w:
-        thread_context.wkb_w = WKBWriter()
+        thread_context.wkb_w = WKBWriter(dim=dim)
+    else:
         thread_context.wkb_w.outdim = dim
     return thread_context.wkb_w
 
 
 def ewkb_w(dim=2):
     if not thread_context.ewkb_w:
-        thread_context.ewkb_w = WKBWriter()
+        thread_context.ewkb_w = WKBWriter(dim=dim)
         thread_context.ewkb_w.srid = True
+    else:
         thread_context.ewkb_w.outdim = dim
     return thread_context.ewkb_w

docs/ref/contrib/gis/geos.txt

@@ -957,12 +957,18 @@ WKB or WKT of the given geometry.  In addition, :class:`WKBWriter` objects
 also have properties that may be used to change the byte order, and or
 include the SRID value (in other words, EWKB).
 
-.. class:: WKBWriter
+.. class:: WKBWriter(dim=2)
 
 ``WKBWriter`` provides the most control over its output.  By default it
 returns OGC-compliant WKB when its ``write`` method is called.  However,
 it has properties that allow for the creation of EWKB, a superset of the
-WKB standard that includes additional information.
+WKB standard that includes additional information. See the
+:attr:`WKBWriter.outdim` documentation for more details about the ``dim``
+argument.
+
+.. versionchanged:: 1.10
+
+    The ability to pass the ``dim`` argument to the constructor was added.
 
 .. method:: WKBWriter.write(geom)
 
@@ -1047,7 +1053,16 @@ geometry should be included with the WKB representation.  Example::
     >>> wkb_w.write_hex(pnt)
     '0101000020E6100000000000000000F03F000000000000F03F'
 
-.. class:: WKTWriter
+.. class:: WKTWriter(dim=2, trim=False, precision=None)
+
+    This class allows outputting the WKT representation of a geometry. See the
+    :attr:`WKBWriter.outdim`, :attr:`trim`, and :attr:`precision` attributes for
+    details about the constructor arguments.
+
+    .. versionchanged:: 1.10
+
+        The ability to pass the ``dim``, ``trim``, and ``precision`` arguments
+        to the constructor was added.
 
 .. method:: WKTWriter.write(geom)
 
@@ -1059,6 +1074,10 @@ Returns the WKT of the given geometry. Example::
     >>> wkt_w.write(pnt)
     'POINT (1.0000000000000000 1.0000000000000000)'
 
+.. attribute:: WKTWriter.outdim
+
+    See :attr:`WKBWriter.outdim`.
+
 .. attribute:: WKTWriter.trim
 
 .. versionadded:: 1.10

tests/gis_tests/geos_tests/test_io.py

@@ -42,6 +42,12 @@ class GEOSIOTest(SimpleTestCase):
         ref_wkt = 'POINT (5.0000000000000000 23.0000000000000000)'
         self.assertEqual(ref_wkt, wkt_w.write(ref).decode())
 
+    def test_wktwriter_constructor_arguments(self):
+        wkt_w = WKTWriter(dim=3, trim=True, precision=3)
+        ref = GEOSGeometry('POINT (5.34562 23 1.5)')
+        ref_wkt = 'POINT Z (5.35 23 1.5)'
+        self.assertEqual(ref_wkt, wkt_w.write(ref).decode())
+
     def test03_wkbreader(self):
         # Creating a WKBReader instance
         wkb_r = WKBReader()
@@ -101,7 +107,7 @@ class GEOSIOTest(SimpleTestCase):
 
         # Ensuring bad output dimensions are not accepted
         for bad_outdim in (-1, 0, 1, 4, 423, 'foo', None):
-            with self.assertRaises(ValueError):
+            with self.assertRaisesMessage(ValueError, 'WKB output dimension must be 2 or 3'):
                 wkb_w.outdim = bad_outdim
 
         # Now setting the output dimensions to be 3