Merge pull request #103 from eseglem/feature/add-has-z-function

Add has_z function to Geometries.

Author: vincentsarago

geojson_pydantic/geometries.py

@@ -23,16 +23,32 @@ def _position_wkt_coordinates(position: Position) -> str:
     return " ".join(str(number) for number in position)
 
 
+def _position_has_z(position: Position) -> bool:
+    return len(position) == 3
+
+
 def _position_list_wkt_coordinates(positions: List[Position]) -> str:
     """Converts a list of Positions to WKT Coordinates."""
     return ", ".join(_position_wkt_coordinates(position) for position in positions)
 
 
+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:
     """Converts lines to WKT Coordinates."""
     return ", ".join(f"({_position_list_wkt_coordinates(line)})" for line in lines)
 
 
+def _lines_has_z(lines: List[List[Position]]) -> 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
+    )
+
+
 class _GeometryBase(BaseModel, abc.ABC):
     """Base class for geometry models"""
 
@@ -49,13 +65,13 @@ def __geo_interface__(self) -> Dict[str, Any]:
 
     @property
     @abc.abstractmethod
-    def _wkt_coordinates(self) -> str:
+    def has_z(self) -> bool:
+        """Checks if any coordinate has a Z value."""
         ...
 
     @property
     @abc.abstractmethod
-    def _wkt_inset(self) -> str:
-        """Return Z for 3 dimensional geometry or an empty string for 2 dimensions."""
+    def _wkt_coordinates(self) -> str:
         ...
 
     @property
@@ -66,11 +82,18 @@ def _wkt_type(self) -> str:
     @property
     def wkt(self) -> str:
         """Return the Well Known Text representation."""
-        return self._wkt_type + (
-            f"{self._wkt_inset}({self._wkt_coordinates})"
-            if self.coordinates
-            else " EMPTY"
-        )
+        # Start with the WKT Type
+        wkt = self._wkt_type
+        if self.coordinates:
+            # If any of the coordinates have a Z add a "Z" to the WKT
+            wkt += " Z " if self.has_z else " "
+            # Add the rest of the WKT inside parentheses
+            wkt += f"({self._wkt_coordinates})"
+        else:
+            # Otherwise it will be "EMPTY"
+            wkt += " EMPTY"
+
+        return wkt
 
 
 class Point(_GeometryBase):
@@ -80,12 +103,13 @@ class Point(_GeometryBase):
     coordinates: Position
 
     @property
-    def _wkt_coordinates(self) -> str:
-        return _position_wkt_coordinates(self.coordinates)
+    def has_z(self) -> bool:
+        """Checks if any coordinate has a Z value."""
+        return _position_has_z(self.coordinates)
 
     @property
-    def _wkt_inset(self) -> str:
-        return " Z " if len(self.coordinates) == 3 else " "
+    def _wkt_coordinates(self) -> str:
+        return _position_wkt_coordinates(self.coordinates)
 
 
 class MultiPoint(_GeometryBase):
@@ -95,8 +119,9 @@ class MultiPoint(_GeometryBase):
     coordinates: MultiPointCoords
 
     @property
-    def _wkt_inset(self) -> str:
-        return " Z " if len(self.coordinates[0]) == 3 else " "
+    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:
@@ -110,8 +135,9 @@ class LineString(_GeometryBase):
     coordinates: LineStringCoords
 
     @property
-    def _wkt_inset(self) -> str:
-        return " Z " if len(self.coordinates[0]) == 3 else " "
+    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:
@@ -125,8 +151,9 @@ class MultiLineString(_GeometryBase):
     coordinates: MultiLineStringCoords
 
     @property
-    def _wkt_inset(self) -> str:
-        return " Z " if len(self.coordinates[0][0]) == 3 else " "
+    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:
@@ -172,8 +199,9 @@ def interiors(self) -> Iterator[LinearRing]:
         )
 
     @property
-    def _wkt_inset(self) -> str:
-        return " Z " if len(self.coordinates[0][0]) == 3 else " "
+    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:
@@ -199,8 +227,9 @@ class MultiPolygon(_GeometryBase):
     coordinates: MultiPolygonCoords
 
     @property
-    def _wkt_inset(self) -> str:
-        return " Z " if len(self.coordinates[0][0][0]) == 3 else " "
+    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:

tests/test_geometries.py

@@ -472,6 +472,125 @@ class PointType(Point):
     )
 
 
+@pytest.mark.parametrize(
+    "coordinates,expected",
+    [
+        ((0, 0), False),
+        ((0, 0, 0), True),
+    ],
+)
+def test_point_has_z(coordinates, expected):
+    assert Point(type="Point", coordinates=coordinates).has_z == expected
+
+
+@pytest.mark.parametrize(
+    "coordinates,expected",
+    [
+        ([(0, 0)], False),
+        ([(0, 0), (1, 1)], False),
+        ([(0, 0), (1, 1, 1)], True),
+        ([(0, 0, 0)], True),
+        ([(0, 0, 0), (1, 1)], True),
+    ],
+)
+def test_multipoint_has_z(coordinates, expected):
+    assert MultiPoint(type="MultiPoint", coordinates=coordinates).has_z == expected
+
+
+@pytest.mark.parametrize(
+    "coordinates,expected",
+    [
+        ([(0, 0), (1, 1)], False),
+        ([(0, 0), (1, 1, 1)], True),
+        ([(0, 0, 0), (1, 1, 1)], True),
+        ([(0, 0, 0), (1, 1)], True),
+    ],
+)
+def test_linestring_has_z(coordinates, expected):
+    assert LineString(type="LineString", coordinates=coordinates).has_z == expected
+
+
+@pytest.mark.parametrize(
+    "coordinates,expected",
+    [
+        ([[(0, 0), (1, 1)]], False),
+        ([[(0, 0), (1, 1)], [(0, 0), (1, 1)]], False),
+        ([[(0, 0), (1, 1)], [(0, 0, 0), (1, 1)]], True),
+        ([[(0, 0), (1, 1)], [(0, 0), (1, 1, 1)]], True),
+        ([[(0, 0), (1, 1, 1)]], True),
+        ([[(0, 0, 0), (1, 1, 1)]], True),
+        ([[(0, 0, 0), (1, 1)]], True),
+        ([[(0, 0, 0), (1, 1, 1)], [(0, 0, 0), (1, 1, 1)]], True),
+    ],
+)
+def test_multilinestring_has_z(coordinates, expected):
+    assert (
+        MultiLineString(type="MultiLineString", coordinates=coordinates).has_z
+        == expected
+    )
+
+
+@pytest.mark.parametrize(
+    "coordinates,expected",
+    [
+        ([[(0, 0), (1, 1), (2, 2), (0, 0)]], False),
+        ([[(0, 0), (1, 1), (2, 2, 2), (0, 0)]], True),
+        ([[(0, 0), (1, 1), (2, 2), (0, 0)], [(0, 0), (1, 1), (2, 2), (0, 0)]], False),
+        (
+            [[(0, 0), (1, 1), (2, 2), (0, 0)], [(0, 0), (1, 1), (2, 2, 2), (0, 0)]],
+            True,
+        ),
+        ([[(0, 0, 0), (1, 1, 1), (2, 2, 2), (0, 0, 0)]], True),
+        (
+            [
+                [(0, 0, 0), (1, 1, 1), (2, 2, 2), (0, 0, 0)],
+                [(0, 0), (1, 1), (2, 2), (0, 0)],
+            ],
+            True,
+        ),
+    ],
+)
+def test_polygon_has_z(coordinates, expected):
+    assert Polygon(type="Polygon", coordinates=coordinates).has_z == expected
+
+
+@pytest.mark.parametrize(
+    "coordinates,expected",
+    [
+        ([[[(0, 0), (1, 1), (2, 2), (0, 0)]]], False),
+        ([[[(0, 0), (1, 1), (2, 2, 2), (0, 0)]]], True),
+        (
+            [[[(0, 0), (1, 1), (2, 2), (0, 0)]], [[(0, 0), (1, 1), (2, 2), (0, 0)]]],
+            False,
+        ),
+        (
+            [
+                [[(0, 0), (1, 1), (2, 2), (0, 0)]],
+                [
+                    [(0, 0), (1, 1), (2, 2), (0, 0)],
+                    [(0, 0, 0), (1, 1, 1), (2, 2, 2), (0, 0, 0)],
+                ],
+            ],
+            True,
+        ),
+        (
+            [[[(0, 0), (1, 1), (2, 2), (0, 0)]], [[(0, 0), (1, 1), (2, 2, 2), (0, 0)]]],
+            True,
+        ),
+        ([[[(0, 0, 0), (1, 1, 1), (2, 2, 2), (0, 0, 0)]]], True),
+        (
+            [
+                [[(0, 0, 0), (1, 1, 1), (2, 2, 2), (0, 0, 0)]],
+                [[(0, 0), (1, 1), (2, 2), (0, 0)]],
+            ],
+            True,
+        ),
+    ],
+)
+def test_multipolygon_has_z(coordinates, expected):
+    assert MultiPolygon(type="MultiPolygon", coordinates=coordinates).has_z == expected
+
+
 @pytest.mark.parametrize(
     "shape",
     [