Add GeoPandas adaptor

Author: nathanjmcdougall

pins/_adaptors.py

@@ -8,17 +8,23 @@
 from typing_extensions import TypeAlias
 
 if TYPE_CHECKING:
+    import geopandas as gpd
     import pandas as pd
 
     PandasDataFrame: TypeAlias = pd.DataFrame
-    DataFrame: TypeAlias = PandasDataFrame
+    GeoPandasGeoDataFrame: TypeAlias = gpd.GeoDataFrame
+    DataFrame: TypeAlias = PandasDataFrame | GeoPandasGeoDataFrame
 
 
 class AbstractPandasFrame(AbstractBackend):
     _backends = [("pandas", "DataFrame")]
 
 
-AbstractDF: TypeAlias = AbstractPandasFrame
+class AbstractGeoPandasFrame(AbstractPandasFrame):
+    _backends = [("geopandas", "GeoDataFrame")]
+
+
+AbstractDF: TypeAlias = AbstractPandasFrame | AbstractGeoPandasFrame
 
 
 class Adaptor:
@@ -142,12 +148,29 @@ def write_feather(self, file: str) -> None:
         self._d.to_feather(file)
 
 
+class GeoPandasAdaptor(PandasAdaptor):
+    _d: ClassVar[GeoPandasGeoDataFrame]  # type: ignore[reportIncompatibleVariableOverride]
+
+    def __init__(self, data: AbstractGeoPandasFrame) -> None:
+        super().__init__(data)
+
+    @property
+    def df_type(self) -> str:
+        # Consider over-riding this for specialized dataframes
+        return "GeoDataFrame"
+
+    def head(self, n: int) -> GeoPandasAdaptor:
+        return GeoPandasAdaptor(self._d.head(n))
+
+
 @overload
 def create_adaptor(obj: DataFrame) -> DFAdaptor: ...
 @overload
 def create_adaptor(obj: Any) -> Adaptor: ...
 def create_adaptor(obj: Any | DataFrame) -> Adaptor | DFAdaptor:
-    if isinstance(obj, AbstractPandasFrame):
+    if isinstance(obj, AbstractGeoPandasFrame):
+        return GeoPandasAdaptor(obj)
+    elif isinstance(obj, AbstractPandasFrame):
         return PandasAdaptor(obj)
     elif isinstance(obj, Adaptor):
         return obj

pins/drivers.py

@@ -174,6 +174,8 @@ def save_data(
         raise NotImplementedError(msg)
     elif pin_type == "parquet":
         adaptor.write_parquet(final_name)
+    elif pin_type == "geoparquet":
+        adaptor.write_parquet(final_name)
     elif pin_type == "joblib":
         adaptor.write_joblib(final_name)
     elif pin_type == "json":