Improve schema

Author: sjpfenninger

workflow/scripts/_schemas.py

@@ -1,25 +1,28 @@
 """Schemas for tabular data used in the workflow."""
 
-from pandera.pandas import DataFrameModel, Field, check
+import pandera.pandas as pa
 from pandera.typing.geopandas import GeoSeries
 from pandera.typing.pandas import Series
-from shapely.geometry import Point
 
 
-class Shapes(DataFrameModel):
-    class Config:
-        coerce = True
-        strict = False
+class ShapesSchema(pa.DataFrameModel):
+    """Schema for geographic shapes."""
 
-    shape_id: Series[str] = Field(unique=True)
-    "Unique ID for this shape."
+    shape_id: Series[str] = pa.Field(unique=True)
+    "A unique identifier for this shape."
     country_id: Series[str]
-    "ISO alpha-3 code."
-    shape_class: Series[str] = Field(isin=["land", "maritime"])
-    "Shape classifier"
-    geometry: GeoSeries[Point] = Field()
-    "Shape polygon."
+    "Country ISO alpha-3 code."
+    shape_class: Series[str] = pa.Field(isin=["land", "maritime"])
+    "Identifier of the shape's context."
+    geometry: GeoSeries
+    "Shape (multi)polygon."
+    parent_name: Series[str] | None
+    "Human-readable name in the parent dataset."
 
-    @check("geometry", element_wise=True)
-    def geom_not_empty(cls, geom):
+    @pa.check("geometry", element_wise=True)
+    def check_geometries(cls, geom):
         return (geom is not None) and (not geom.is_empty) and geom.is_valid
+
+    class Config:
+        coerce = True
+        strict = False

workflow/scripts/breakup_shape.py

@@ -4,7 +4,7 @@
 
 import click
 import geopandas as gpd
-from _schemas import Shapes
+from _schemas import ShapesSchema
 
 
 @click.command()
@@ -26,7 +26,7 @@ def breakup_shape(shapes_path, split_by, output_path):
     output_path = Path(output_path)
     output_path.mkdir(parents=True, exist_ok=True)
     shapes = gpd.read_parquet(shapes_path)
-    shapes = Shapes.validate(shapes)
+    shapes = ShapesSchema.validate(shapes)
 
     # Print rows where geometry is empty
     if shapes.geometry.is_empty.any():