Repository: ecmwf/earthkit-data
Pull request #159 (2023-08-23 15:54:42)

Feature/pandas geopandas transformers
pandas/geopandas transformers added, with tests


Patch:
diff --git a/earthkit/data/readers/geojson.py b/earthkit/data/readers/geojson.py
index 60f83dd..81aa80c 100644
--- a/earthkit/data/readers/geojson.py
+++ b/earthkit/data/readers/geojson.py
@@ -13,7 +13,7 @@ import mimetypes
 
 import numpy as np
 
-from . import Reader
+from earthkit.data.readers import Reader
 
 
 class GeojsonReader(Reader):
@@ -29,7 +29,16 @@ class GeojsonReader(Reader):
             self.fields = self.get_fields()
 
     def __repr__(self):
-        return "GeojsonReader(%s)" % (self.path,)
+        return f"GeojsonReader({self.path})"
+        # return self.to_pandas().__repr__()
+
+    def _repr_html_(self):
+        html_repr = (
+            f"<h3>GeojsonReader(represented as a geopandas object):</h3>"
+            f"{self.to_pandas()._repr_html_()}"
+        )
+
+        return html_repr
 
     def __iter__(self):
         """
@@ -46,15 +55,18 @@ class GeojsonReader(Reader):
         self._scan()
         return self.fields[n]
 
+    def mutate_source(self):
+        # A Geojson is a source itself
+        return self
+
     def bounding_box(self, **kwargs):
         return self.to_pandas(**kwargs).crs.area_of_use.bounds
 
     def ls(self, **kwargs):
         return self.to_pandas(**kwargs)
 
-    describe = ls
-    # def describe(self, kwargs):
-    #     return self.to_pandas(**kwargs)
+    def describe(self, **kwargs):
+        return self.to_pandas(**kwargs)
 
     def get_fields(self, **kwargs):
         return [row[1] for row in self.to_pandas(**kwargs).iterrows()]
@@ -65,15 +77,13 @@ class GeojsonReader(Reader):
             arr = arr.flatten()
         return arr
 
-    # def to_geopandas(self, **kwargs):
-    #     # TODO: handle multiple paths
-    #     return self.to_pandas(**kwargs)
-
     def to_pandas(self, **kwargs):
         # TODO: handle multiple paths
         return self.to_pandas_from_multi_paths([self.path], **kwargs)
 
-    to_geopandas = to_pandas
+    def to_geopandas(self, **kwargs):
+        # TODO: handle multiple paths
+        return self.to_pandas(**kwargs)
 
     def to_xarray(self, **kwargs):
         return self.to_pandas(**kwargs).to_xarray()
diff --git a/earthkit/data/translators/pandas.py b/earthkit/data/translators/pandas.py
new file mode 100644
index 0000000..9998f8d
--- /dev/null
+++ b/earthkit/data/translators/pandas.py
@@ -0,0 +1,73 @@
+# (C) Copyright 2022 ECMWF.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+import pandas as pd
+
+from earthkit.data.translators import Translator
+
+
+class PandasSeriesTranslator(Translator):
+    """Translator class for pandas `Series`"""
+
+    def __init__(self, data, *args, **kwargs):
+        self.data = data.to_pandas(*args, **kwargs)
+
+    def __call__(self):
+        """
+        Series requested, if DataFrame return the first column
+        """
+        if isinstance(self.data, pd.DataFrame):
+            return self.data.iloc[:, 0]
+
+        return self.data
+
+
+class PandasDataFrameTranslator(PandasSeriesTranslator):
+    """Translator class for pandas `DataFrame`"""
+
+    def __call__(self):
+        """
+        DataFrame requested, if Series convert to DataFrame.
+        """
+        if isinstance(self.data, pd.Series):
+            return self.data.to_frame()
+
+        return self.data
+
+
+class GeoPandasDataFrameTranslator(PandasSeriesTranslator):
+    """Translator class for geopandas `DataFrame`"""
+
+    def __call__(self):
+        """
+        GeoDataFrame requested, if normal pandas convert to geopandas.
+        """
+        import geopandas as gpd
+
+        if isinstance(self.data, pd.DataFrame):
+            return gpd.GeoDataFrame(self.data)
+
+        return self.data
+
+
+def translator(data, cls, *args, **kwargs):
+    try:
+        import geopandas as gpd
+
+        l_gpd = True
+    except ImportError:
+        l_gpd = False
+
+    if l_gpd and cls in [gpd.geodataframe.GeoDataFrame]:
+        return GeoPandasDataFrameTranslator(data, *args, **kwargs)
+    if cls in [pd.DataFrame]:
+        return PandasDataFrameTranslator(data, *args, **kwargs)
+    if cls in [pd.Series]:
+        return PandasSeriesTranslator(data, *args, **kwargs)
+
+    return None
diff --git a/earthkit/data/wrappers/pandas.py b/earthkit/data/wrappers/pandas.py
index ac4a3e3..0c58c93 100644
--- a/earthkit/data/wrappers/pandas.py
+++ b/earthkit/data/wrappers/pandas.py
@@ -165,7 +165,7 @@ def wrapper(data, *args, **kwargs):
         return GeoPandasDataFrameWrapper(data, *args, **kwargs)
     if isinstance(data, pd.DataFrame):
         return PandasDataFrameWrapper(data, *args, **kwargs)
-    elif isinstance(data, pd.Series):
+    if isinstance(data, pd.Series):
         return PandasSeriesWrapper(data, *args, **kwargs)
 
     return None
diff --git a/pyproject.toml b/pyproject.toml
index d1c115a..cc8bf45 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -9,7 +9,7 @@ profile = "black"
 
 [tool.pydocstyle]
 convention = "numpy"
-add_ignore = ["D1", "D200", "D205", "D400", "D401"]
+add_ignore = ["D1", "D200", "D205", "D400", "D401", "D403"]
 
 [tool.setuptools_scm]
 write_to = "earthkit/data/version.py"
diff --git a/setup.cfg b/setup.cfg
index 3f02006..f909315 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -29,8 +29,8 @@ install_requires =
     dask
     entrypoints
     filelock
-    geopandas
     jinja2
+    markdown
     multiurl
     netcdf4
     pdbufr>=0.11.0


Test patch:
diff --git a/tests/translators/test_translators.py b/tests/translators/test_translators.py
index 36a868e..f4a9250 100644
--- a/tests/translators/test_translators.py
+++ b/tests/translators/test_translators.py
@@ -12,11 +12,14 @@
 
 import logging
 
+import geopandas as gpd
 import numpy as np
+import pandas as pd
 import xarray as xr
 
 from earthkit.data import from_source, transform, translators, wrappers
 from earthkit.data.translators import ndarray as ndtranslator
+from earthkit.data.translators import pandas as pdtranslator
 from earthkit.data.translators import string as strtranslator
 from earthkit.data.translators import xarray as xrtranslator
 
@@ -72,17 +75,62 @@ def test_xr_dataset_translator():
     # Check that an xr.Dataset translator can be created
     _xrwrapper = wrappers.get_wrapper(xr.Dataset())
     _trans = xrtranslator.translator(_xrwrapper, xr.Dataset)
-    assert isinstance(_trans, xrtranslator.XArrayDataArrayTranslator)
+    assert isinstance(_trans, xrtranslator.XArrayDatasetTranslator)
 
     # Check that get_translator method find the right translator
     _trans = translators.get_translator(xr.Dataset(), xr.Dataset)
-    assert isinstance(_trans, xrtranslator.XArrayDataArrayTranslator)
+    assert isinstance(_trans, xrtranslator.XArrayDatasetTranslator)
 
     # Check that public API method transforms to correct type (back to original)
     transformed = transform(xr.Dataset(), xr.Dataset)
     assert isinstance(transformed, xr.Dataset)
 
 
+def test_pd_series_translator():
+    # Check that an xr.Dataset translator can be created
+    _pdwrapper = wrappers.get_wrapper(pd.Series())
+    _trans = pdtranslator.translator(_pdwrapper, pd.DataFrame)
+    assert isinstance(_trans, pdtranslator.PandasSeriesTranslator)
+
+    # Check that get_translator method find the right translator
+    _trans = translators.get_translator(pd.Series(), pd.Series)
+    assert isinstance(_trans, pdtranslator.PandasSeriesTranslator)
+
+    # Check that public API method transforms to correct type (back to original)
+    transformed = transform(pd.Series(), pd.Series)
+    assert isinstance(transformed, pd.Series)
+
+
+def test_pd_dataframe_translator():
+    # Check that an xr.Dataset translator can be created
+    _pdwrapper = wrappers.get_wrapper(pd.DataFrame())
+    _trans = pdtranslator.translator(_pdwrapper, pd.DataFrame)
+    assert isinstance(_trans, pdtranslator.PandasDataFrameTranslator)
+
+    # Check that get_translator method find the right translator
+    _trans = translators.get_translator(pd.DataFrame(), pd.DataFrame)
+    assert isinstance(_trans, pdtranslator.PandasDataFrameTranslator)
+
+    # Check that public API method transforms to correct type (back to original)
+    transformed = transform(pd.DataFrame(), pd.DataFrame)
+    assert isinstance(transformed, pd.DataFrame)
+
+
+def test_gpd_dataframe_translator():
+    # Check that an xr.Dataset translator can be created
+    _pdwrapper = wrappers.get_wrapper(gpd.GeoDataFrame())
+    _trans = pdtranslator.translator(_pdwrapper, gpd.GeoDataFrame)
+    assert isinstance(_trans, pdtranslator.GeoPandasDataFrameTranslator)
+
+    # Check that get_translator method find the right translator
+    _trans = translators.get_translator(gpd.GeoDataFrame(), gpd.GeoDataFrame)
+    assert isinstance(_trans, pdtranslator.GeoPandasDataFrameTranslator)
+
+    # Check that public API method transforms to correct type (back to original)
+    transformed = transform(gpd.GeoDataFrame(), gpd.GeoDataFrame)
+    assert isinstance(transformed, gpd.GeoDataFrame)
+
+
 def test_transform_from_grib_file():
     # transform grib-based data object
     f = from_source("file", "tests/data/test_single.grib")