Dataset is the exception case, not DataArray

ecmwf · Aug 25, 2023 · 539ef64 · 539ef64
1 parent 47b6b6c
commit 539ef64
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 13 deletions.
diff --git a/ci/combined-environment-ci.yml b/ci/combined-environment-ci.yml
@@ -16,3 +16,6 @@ dependencies:
 - sphinx
 - sphinx-autoapi
 - xarray
+- pip:
+  - git+https://github.com/ecmwf/multiurl
+
diff --git a/earthkit/climate/aggregate.py b/earthkit/climate/aggregate.py
@@ -259,15 +259,15 @@ def _reduce_dataarray(
 
     else:
         # If how is string, fetch function from dictionary:
-        if isinstance(how, str):
-            if how in dir(dataarray):
-                red_array = dataarray.__getattribute__(how)(**kwargs)
-            else:
+        if isinstance(how, str) and how in dir(dataarray):
+            red_array = dataarray.__getattribute__(how)(**kwargs)
+        else:
+            if isinstance(how, str):
                 how_label = deepcopy(how)
                 how = get_how(how)
-        assert isinstance(how, T.Callable), f"how method not recognised: {how}"
+            assert isinstance(how, T.Callable), f"how method not recognised: {how}"
 
-        red_array = dataarray.reduce(how, **kwargs)
+            red_array = dataarray.reduce(how, **kwargs)
 
     if how_label:
         red_array = red_array.rename(f"{red_array.name}_{how_label}")
@@ -313,12 +313,12 @@ def reduce(
         A data array with reduce dimensions removed.
 
     """
-    if isinstance(dataarray, xr.DataArray):
-        return _reduce_dataarray(dataarray, **kwargs)
-    else:
+    if isinstance(dataarray, (xr.Dataset)):
         return xr.Dataset(
             [_reduce_dataarray(dataarray[var], **kwargs) for var in dataarray.data_vars]
         )
+    else:
+        return _reduce_dataarray(dataarray, **kwargs)
 
 
 def rolling_reduce(

diff --git a/earthkit/climate/shapes.py b/earthkit/climate/shapes.py
@@ -269,10 +269,7 @@ def reduce(
         Each slice of layer corresponds to a feature in layer.
 
     """
-    # kwargs.update({"how": how})
-    if isinstance(dataarray, xr.DataArray):
-        return _reduce_dataarray(dataarray, geodataframe, **kwargs)
-    else:
+    if isinstance(dataarray, xr.Dataset):
         if kwargs.get("return_as", "pandas") in ["xarray"]:
             return xr.Dataset(
                 [
@@ -285,6 +282,9 @@ def reduce(
             for var in dataarray.data_vars:
                 out = _reduce_dataarray(dataarray[var], geodataframe, **kwargs)
             return out
+    else:
+        return _reduce_dataarray(dataarray, geodataframe, **kwargs)
+
 
 
 def _reduce_dataarray(