From ed740afffa6df06e13f3e08bcfb1036c337788bc Mon Sep 17 00:00:00 2001
From: Panos Mavrogiorgos <pmav99@gmail.com>
Date: Thu, 22 Aug 2024 11:49:09 +0300
Subject: [PATCH] feat: Add `to_geodataframe()`

---
 pyposeidon/boundary.py |  7 ++++---
 pyposeidon/tools.py    | 11 +++++++++++
 tests/test_boundary.py |  7 ++++---
 tests/test_dem_fix.py  |  9 ++++-----
 4 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/pyposeidon/boundary.py b/pyposeidon/boundary.py
index 7ba181e..966480e 100644
--- a/pyposeidon/boundary.py
+++ b/pyposeidon/boundary.py
@@ -16,6 +16,7 @@
 import logging
 import shapely
 from tqdm.auto import tqdm
+from pyposeidon.tools import to_geodataframe
 from pyposeidon.utils.coastfix import simplify
 import sys
 
@@ -56,7 +57,7 @@ def __init__(self, **kwargs):
 
         elif isinstance(coastlines, str):
             logger.info("reading {}".format(coastlines))
-            coasts = gp.GeoDataFrame.from_file(coastlines)
+            coasts = to_geodataframe(coastlines)
             # check coastlines
             if coasts.buffer(0).is_valid.all() and (coasts.buffer(0).boundary.geom_type == "LineString").all():
                 self.coasts = gp.GeoDataFrame(geometry=coasts.buffer(0))
@@ -96,7 +97,7 @@ def __init__(self, **kwargs):
 
             else:
                 try:
-                    self.geometry = gp.read_file(geometry)
+                    self.geometry = to_geodataframe(geometry)
                 except:
                     logger.warning("geometry is not a file, trying with geopandas Dataset")
                     if isinstance(geometry, gp.GeoDataFrame):
@@ -107,7 +108,7 @@ def __init__(self, **kwargs):
 
         else:
             try:
-                self.geometry = gp.read_file(geometry)
+                self.geometry = to_geodataframe(geometry)
             except:
                 logger.warning("geometry is not a file, trying with geopandas Dataset")
                 if isinstance(geometry, gp.GeoDataFrame):
diff --git a/pyposeidon/tools.py b/pyposeidon/tools.py
index ec91e77..095ed61 100644
--- a/pyposeidon/tools.py
+++ b/pyposeidon/tools.py
@@ -583,3 +583,14 @@ def get_netcdf_encoding(
             )
             update_or_add(encoding, var, params)
     return encoding
+
+
+def to_geodataframe(
+    path: str | os.PathLike[str],
+    **kwargs: T.Any,
+) -> gpd.GeoDataFrame:
+    if str(path).endswith("parquet") or str(path).endswith("pq"):
+        gdf = gpd.read_parquet(path=path, **kwargs)
+    else:
+        gdf = gpd.read_file(filename=path, **kwargs)
+    return gdf
diff --git a/tests/test_boundary.py b/tests/test_boundary.py
index 7d19d75..c0c7795 100644
--- a/tests/test_boundary.py
+++ b/tests/test_boundary.py
@@ -1,14 +1,15 @@
-import pyposeidon.boundary as pb
-import pytest
 import geopandas as gp
+import pytest
 
+import pyposeidon.boundary as pb
 from . import DATA_DIR
+from pyposeidon.tools import to_geodataframe
 
 noaa = DATA_DIR / "bl.zip"
 
 COAST_FILE = (DATA_DIR / "ocean.parquet").as_posix()
 
-land = gp.read_file(COAST_FILE)
+land = to_geodataframe(COAST_FILE)
 coast = gp.GeoDataFrame(geometry=land.boundary)
 
 INPUTS = pytest.mark.parametrize("input", [land, coast])
diff --git a/tests/test_dem_fix.py b/tests/test_dem_fix.py
index 30bdbb5..8b4bbc0 100644
--- a/tests/test_dem_fix.py
+++ b/tests/test_dem_fix.py
@@ -1,10 +1,9 @@
-import pyposeidon.dem as pdem
-import pyposeidon.mesh as pmesh
 import pytest
-import numpy as np
-import geopandas as gp
 
+import pyposeidon.dem as pdem
+import pyposeidon.mesh as pmesh
 from . import DATA_DIR
+from pyposeidon.tools import to_geodataframe
 
 COAST_FILE = (DATA_DIR / "ocean.parquet").as_posix()
 
@@ -29,7 +28,7 @@
 
 @pytest.fixture(scope="session")
 def coasts():
-    coast = gp.read_file(COAST_FILE)
+    coast = to_geodataframe(COAST_FILE)
     return coast