Merge pull request #107 from eseglem/bugfix/102-wkt-mixed-dimensionality

Address mixed dimensionality in wkt

Author: vincentsarago

geojson_pydantic/geometries.py

@@ -1,7 +1,6 @@
 """pydantic models for GeoJSON Geometry objects."""
-
 import abc
-from typing import Any, Iterator, List, Literal, Optional, Union
+from typing import Any, Iterator, List, Literal, Optional, Protocol, Union
 
 from pydantic import BaseModel, Field, ValidationError, validator
 from pydantic.error_wrappers import ErrorWrapper
@@ -20,70 +19,87 @@
 )
 
 
-def _position_wkt_coordinates(position: Position) -> str:
+def _position_wkt_coordinates(coordinates: Position, force_z: bool = False) -> str:
     """Converts a Position to WKT Coordinates."""
-    return " ".join(str(number) for number in position)
+    wkt_coordinates = " ".join(str(number) for number in coordinates)
+    if force_z and len(coordinates) < 3:
+        wkt_coordinates += " 0.0"
+    return wkt_coordinates
 
 
 def _position_has_z(position: Position) -> bool:
     return len(position) == 3
 
 
-def _position_list_wkt_coordinates(positions: List[Position]) -> str:
+def _position_list_wkt_coordinates(
+    coordinates: List[Position], force_z: bool = False
+) -> str:
     """Converts a list of Positions to WKT Coordinates."""
-    return ", ".join(_position_wkt_coordinates(position) for position in positions)
+    return ", ".join(
+        _position_wkt_coordinates(position, force_z) for position in coordinates
+    )
 
 
 def _position_list_has_z(positions: List[Position]) -> bool:
     """Checks if any position in a list has a Z."""
     return any(_position_has_z(position) for position in positions)
 
 
-def _lines_wtk_coordinates(lines: List[List[Position]]) -> str:
+def _lines_wtk_coordinates(
+    coordinates: List[LineStringCoords], force_z: bool = False
+) -> str:
     """Converts lines to WKT Coordinates."""
-    return ", ".join(f"({_position_list_wkt_coordinates(line)})" for line in lines)
+    return ", ".join(
+        f"({_position_list_wkt_coordinates(line, force_z)})" for line in coordinates
+    )
 
 
-def _lines_has_z(lines: List[List[Position]]) -> bool:
+def _lines_has_z(lines: List[LineStringCoords]) -> bool:
     """Checks if any position in a list has a Z."""
     return any(
         _position_has_z(position) for positions in lines for position in positions
     )
 
 
+def _polygons_wkt_coordinates(
+    coordinates: List[PolygonCoords], force_z: bool = False
+) -> str:
+    return ",".join(
+        f"({_lines_wtk_coordinates(polygon, force_z)})" for polygon in coordinates
+    )
+
+
+class _WktCallable(Protocol):
+    def __call__(self, coordinates: Any, force_z: bool) -> str:
+        ...
+
+
 class _GeometryBase(BaseModel, abc.ABC, GeoInterfaceMixin):
     """Base class for geometry models"""
 
     type: str
     coordinates: Any
     bbox: Optional[BBox] = None
 
+    __wkt_coordinates__: _WktCallable
+
     @property
     @abc.abstractmethod
     def has_z(self) -> bool:
         """Checks if any coordinate has a Z value."""
         ...
 
-    @property
-    @abc.abstractmethod
-    def _wkt_coordinates(self) -> str:
-        ...
-
-    @property
-    def _wkt_type(self) -> str:
-        """Return the WKT name of the geometry."""
-        return self.type.upper()
-
     @property
     def wkt(self) -> str:
         """Return the Well Known Text representation."""
         # Start with the WKT Type
-        wkt = self._wkt_type
+        wkt = self.type.upper()
+        has_z = self.has_z
         if self.coordinates:
             # If any of the coordinates have a Z add a "Z" to the WKT
-            wkt += " Z " if self.has_z else " "
+            wkt += " Z " if has_z else " "
             # Add the rest of the WKT inside parentheses
-            wkt += f"({self._wkt_coordinates})"
+            wkt += f"({self.__wkt_coordinates__(self.coordinates, force_z=has_z)})"
         else:
             # Otherwise it will be "EMPTY"
             wkt += " EMPTY"
@@ -97,70 +113,64 @@ class Point(_GeometryBase):
     type: Literal["Point"]
     coordinates: Position
 
+    __wkt_coordinates__ = staticmethod(_position_wkt_coordinates)
+
     @property
     def has_z(self) -> bool:
         """Checks if any coordinate has a Z value."""
         return _position_has_z(self.coordinates)
 
-    @property
-    def _wkt_coordinates(self) -> str:
-        return _position_wkt_coordinates(self.coordinates)
-
 
 class MultiPoint(_GeometryBase):
     """MultiPoint Model"""
 
     type: Literal["MultiPoint"]
     coordinates: MultiPointCoords
 
+    __wkt_coordinates__ = staticmethod(_position_list_wkt_coordinates)
+
     @property
     def has_z(self) -> bool:
         """Checks if any coordinate has a Z value."""
         return _position_list_has_z(self.coordinates)
 
-    @property
-    def _wkt_coordinates(self) -> str:
-        return _position_list_wkt_coordinates(self.coordinates)
-
 
 class LineString(_GeometryBase):
     """LineString Model"""
 
     type: Literal["LineString"]
     coordinates: LineStringCoords
 
+    __wkt_coordinates__ = staticmethod(_position_list_wkt_coordinates)
+
     @property
     def has_z(self) -> bool:
         """Checks if any coordinate has a Z value."""
         return _position_list_has_z(self.coordinates)
 
-    @property
-    def _wkt_coordinates(self) -> str:
-        return _position_list_wkt_coordinates(self.coordinates)
-
 
 class MultiLineString(_GeometryBase):
     """MultiLineString Model"""
 
     type: Literal["MultiLineString"]
     coordinates: MultiLineStringCoords
 
+    __wkt_coordinates__ = staticmethod(_lines_wtk_coordinates)
+
     @property
     def has_z(self) -> bool:
         """Checks if any coordinate has a Z value."""
         return _lines_has_z(self.coordinates)
 
-    @property
-    def _wkt_coordinates(self) -> str:
-        return _lines_wtk_coordinates(self.coordinates)
-
 
 class Polygon(_GeometryBase):
     """Polygon Model"""
 
     type: Literal["Polygon"]
     coordinates: PolygonCoords
 
+    __wkt_coordinates__ = staticmethod(_lines_wtk_coordinates)
+
     @validator("coordinates")
     def check_closure(cls, coordinates: List) -> List:
         """Validate that Polygon is closed (first and last coordinate are the same)."""
@@ -186,10 +196,6 @@ def has_z(self) -> bool:
         """Checks if any coordinates have a Z value."""
         return _lines_has_z(self.coordinates)
 
-    @property
-    def _wkt_coordinates(self) -> str:
-        return _lines_wtk_coordinates(self.coordinates)
-
     @classmethod
     def from_bounds(
         cls, xmin: float, ymin: float, xmax: float, ymax: float
@@ -209,17 +215,13 @@ class MultiPolygon(_GeometryBase):
     type: Literal["MultiPolygon"]
     coordinates: MultiPolygonCoords
 
+    __wkt_coordinates__ = staticmethod(_polygons_wkt_coordinates)
+
     @property
     def has_z(self) -> bool:
         """Checks if any coordinates have a Z value."""
         return any(_lines_has_z(polygon) for polygon in self.coordinates)
 
-    @property
-    def _wkt_coordinates(self) -> str:
-        return ",".join(
-            f"({_lines_wtk_coordinates(polygon)})" for polygon in self.coordinates
-        )
-
     @validator("coordinates")
     def check_closure(cls, coordinates: List) -> List:
         """Validate that Polygon is closed (first and last coordinate are the same)."""
@@ -254,24 +256,23 @@ def __getitem__(self, index: int) -> Geometry:
         """get geometry at a given index"""
         return self.geometries[index]
 
-    @property
-    def _wkt_type(self) -> str:
-        """Return the WKT name of the geometry."""
-        return self.type.upper()
-
-    @property
-    def _wkt_coordinates(self) -> str:
-        """Encode coordinates as WKT."""
-        return ", ".join(geom.wkt for geom in self.geometries)
-
     @property
     def wkt(self) -> str:
         """Return the Well Known Text representation."""
-        return (
-            self._wkt_type
-            + " "
-            + (f"({self._wkt_coordinates})" if self._wkt_coordinates else "EMPTY")
+        # Each geometry will check its own coordinates for Z and include "Z" in the wkt
+        # if necessary. Rather than looking at the coordinates for each of the geometries
+        # again, we can just get the wkt from each of them and check if there is a Z
+        # anywhere in the text.
+
+        # Get the wkt from each of the geometries in the collection
+        geometries = (
+            f'({", ".join(geom.wkt for geom in self.geometries)})'
+            if self.geometries
+            else "EMPTY"
         )
+        # If any of them contain `Z` add Z to the output wkt
+        z = " Z " if "Z" in geometries else " "
+        return f"{self.type.upper()}{z}{geometries}"
 
 
 def parse_geometry_obj(obj: Any) -> Geometry:

tests/test_geometries.py

@@ -59,6 +59,8 @@ def test_point_invalid_coordinates(coordinates):
         [(1.0, 2.0), (1.0, 2.0)],
         # Has Z
         [(1.0, 2.0, 3.0), (1.0, 2.0, 3.0)],
+        # Mixed
+        [(1.0, 2.0), (1.0, 2.0, 3.0)],
     ],
 )
 def test_multi_point_valid_coordinates(coordinates):
@@ -93,6 +95,7 @@ def test_multi_point_invalid_coordinates(coordinates):
         [(1.0, 2.0), (3.0, 4.0), (5.0, 6.0)],
         # Two Points, has Z
         [(0.0, 0.0, 0.0), (1.0, 1.0, 1.0)],
+        # Shapely doesn't like mixed here
     ],
 )
 def test_line_string_valid_coordinates(coordinates):
@@ -130,6 +133,8 @@ def test_line_string_invalid_coordinates(coordinates):
         [[(1.0, 2.0), (3.0, 4.0)], [(0.0, 0.0), (1.0, 1.0)]],
         # Two lines, two points each, has Z
         [[(1.0, 2.0, 0.0), (3.0, 4.0, 1.0)], [(0.0, 0.0, 0.0), (1.0, 1.0, 1.0)]],
+        # Mixed
+        [[(1.0, 2.0), (3.0, 4.0)], [(0.0, 0.0, 0.0), (1.0, 1.0, 1.0)]],
     ],
 )
 def test_multi_line_string_valid_coordinates(coordinates):
@@ -206,6 +211,17 @@ def test_polygon_valid_coordinates(coordinates):
                 (2.0, 2.0, 1.0),
             ],
         ],
+        # Mixed
+        [
+            [(0.0, 0.0), (0.0, 10.0), (10.0, 10.0), (10.0, 0.0), (0.0, 0.0)],
+            [
+                (2.0, 2.0, 2.0),
+                (2.0, 4.0, 0.0),
+                (4.0, 4.0, 0.0),
+                (4.0, 2.0, 0.0),
+                (2.0, 2.0, 2.0),
+            ],
+        ],
     ],
 )
 def test_polygon_with_holes(coordinates):
@@ -271,6 +287,19 @@ def test_polygon_invalid_coordinates(coordinates):
                 ],
             ]
         ],
+        # Mixed
+        [
+            [
+                [(0.0, 0.0), (1.0, 0.0), (1.0, 1.0), (0.0, 1.0), (0.0, 0.0)],
+                [
+                    (2.1, 2.1, 2.1),
+                    (2.2, 2.1, 2.0),
+                    (2.2, 2.2, 2.2),
+                    (2.1, 2.2, 2.3),
+                    (2.1, 2.1, 2.1),
+                ],
+            ]
+        ],
     ],
 )
 def test_multi_polygon(coordinates):
@@ -452,6 +481,18 @@ def test_getitem_geometry_collection(polygon):
     assert item == gc[0]
 
 
+def test_wkt_mixed_geometry_collection():
+    point = Point(type="Point", coordinates=(0.0, 0.0, 0.0))
+    line_string = LineString(type="LineString", coordinates=[(0.0, 0.0), (1.0, 1.0)])
+    gc = GeometryCollection(type="GeometryCollection", geometries=[point, line_string])
+    assert_wkt_equivalence(gc)
+
+
+def test_wkt_empty_geometry_collection():
+    gc = GeometryCollection(type="GeometryCollection", geometries=[])
+    assert_wkt_equivalence(gc)
+
+
 def test_polygon_from_bounds():
     """Result from `from_bounds` class method should be the same."""
     coordinates = [[(1.0, 2.0), (3.0, 2.0), (3.0, 4.0), (1.0, 4.0), (1.0, 2.0)]]