[packages/calitp_data_analysis] Auth helpers for accessing GeoPandas geospatial data in Google Cloud Storage

.pre-commit-config.yaml

@@ -21,7 +21,7 @@ repos:
         # Suppress SyntaxWarning about invalid escape sequence from calitp-data-infra dependency without modifying source
         entry: env PYTHONWARNINGS="ignore::SyntaxWarning" flake8
   - repo: https://github.com/psf/black
-    rev: 23.1.0
+    rev: 24.10.0
     hooks:
       - id: black
         args: ["--config=./pyproject.toml"]

packages/calitp-data-analysis/calitp_data_analysis/gcs_geopandas.py

@@ -0,0 +1,40 @@
+import gcsfs  # type: ignore
+import geopandas  # type: ignore
+import google.auth  # type: ignore
+
+
+class GCSGeoPandas:
+    """
+    GCSGeoPandas contains authentication helpers for interacting with Google Cloud   Storage with GeoPandas
+    """
+
+    def __init__(self):
+        """Sets instance token to credentials returned from Google auth request"""
+        credentials, _ = google.auth.default()
+        self.token = credentials
+
+    def gcs_filesystem(self, **kwargs):
+        """Returns a Google Cloud Storage Filesystem"""
+        return gcsfs.GCSFileSystem(token=self.token, **kwargs)
+
+    def read_parquet(self, path, *args, **kwargs):
+        """Delegates to geopandas.read_parquet with storage option token
+
+        Passes the auth credentials from Google auth as storage option token
+        """
+        storage_options = kwargs.get("storage_options", {}) | {"token": self.token}
+        return geopandas.read_parquet(
+            path, storage_options=storage_options, *args, **kwargs
+        )
+
+    def geo_data_frame_to_parquet(self, geo_data_frame, path, *args, **kwargs):
+        """Delegates to .to_parquet on the passed geo_data_frame, providing Google Cloud Storage file system
+
+        Fetches gcsfs.GCSFileSystem instance
+
+        Passes the auth credentials from Google Cloud Storage as storage option token
+        """
+        gcs_filesystem = self.gcs_filesystem()
+        return geo_data_frame.to_parquet(
+            path, filesystem=gcs_filesystem, *args, **kwargs
+        )