Refactor backend-agnostic lazy-compute utility

src/xregrid/regridder.py

@@ -15,6 +15,7 @@
     is_dask,
     is_cubed,
     _get_min_max_lazy_aware,
+    _compute_lazy_aware,
 )
 from xregrid.constants import (
     get_regrid_method_map,
@@ -2081,18 +2082,8 @@ def _detect_periodicity(self, ds: xr.Dataset) -> bool:
                         "or set the 'boundary' attribute to 'periodic' in your longitude coordinate."
                     )
                     try:
-                        # Use xarray's compute to remain backend-agnostic
-                        # if the tasks are xarray/dask objects.
-                        if hasattr(lon_min_task, "compute"):
-                            lon_min = lon_min_task.compute()
-                        else:
-                            lon_min = lon_min_task
-
-                        if hasattr(lon_max_task, "compute"):
-                            lon_max = lon_max_task.compute()
-                        else:
-                            lon_max = lon_max_task
-
+                        lon_min = _compute_lazy_aware(lon_min_task)
+                        lon_max = _compute_lazy_aware(lon_max_task)
                         lon_min, lon_max = float(lon_min), float(lon_max)
                     except Exception:
                         lon_min = lon_max = None

src/xregrid/utils.py

@@ -33,7 +33,8 @@
 
 def is_cubed(obj: Any) -> bool:
     """
-    Check if an object is a cubed array or a Dataset/DataArray containing one.
+    Check if an object is a cubed array or a Dataset/DataArray containing one,
+    or a collection containing cubed arrays (recursively checks lists, tuples, and dicts).
 
     Parameters
     ----------
@@ -57,12 +58,19 @@ def is_cubed(obj: Any) -> bool:
     if isinstance(obj, xr.Dataset):
         return any(isinstance(v.data, cubed.Array) for v in obj.data_vars.values())
 
+    if isinstance(obj, (list, tuple)):
+        return any(is_cubed(item) for item in obj)
+
+    if isinstance(obj, dict):
+        return any(is_cubed(item) for item in obj.values())
+
     return False
 
 
 def is_dask(obj: Any) -> bool:
     """
-    Check if an object is a dask collection or a Dataset/DataArray containing one.
+    Check if an object is a dask collection or a Dataset/DataArray containing one,
+    or a collection containing dask collections (recursively checks lists, tuples, and dicts).
 
     Parameters
     ----------
@@ -88,6 +96,12 @@ def is_dask(obj: Any) -> bool:
     if isinstance(obj, xr.Dataset):
         return any(is_dask_collection(v.data) for v in obj.data_vars.values())
 
+    if isinstance(obj, (list, tuple)):
+        return any(is_dask(item) for item in obj)
+
+    if isinstance(obj, dict):
+        return any(is_dask(item) for item in obj.values())
+
     return False
 
 
@@ -108,6 +122,74 @@ def is_lazy(obj: Any) -> bool:
     return is_dask(obj) or is_cubed(obj)
 
 
+def _compute_lazy_aware(obj: Any) -> Any:
+    """
+    Centralized backend-agnostic utility for computing lazy data.
+
+    Detects and dispatches to Dask or Cubed compute methods, or falls back to
+    returning eager data. Supports recursive resolution for single objects,
+    lists, tuples, and dictionaries.
+
+    Parameters
+    ----------
+    obj : Any
+        The object (or collection of objects) to compute.
+
+    Returns
+    -------
+    Any
+        The computed eager object(s).
+    """
+    if not is_lazy(obj):
+        if isinstance(obj, list):
+            return [_compute_lazy_aware(item) for item in obj]
+        if isinstance(obj, tuple):
+            return tuple(_compute_lazy_aware(item) for item in obj)
+        if isinstance(obj, dict):
+            return {k: _compute_lazy_aware(v) for k, v in obj.items()}
+        return obj
+
+    if is_cubed(obj):
+        if cubed is None:
+            return obj
+        if isinstance(obj, (list, tuple, dict)):
+            if isinstance(obj, list):
+                return [_compute_lazy_aware(item) for item in obj]
+            if isinstance(obj, tuple):
+                return tuple(_compute_lazy_aware(item) for item in obj)
+            if isinstance(obj, dict):
+                return {k: _compute_lazy_aware(v) for k, v in obj.items()}
+
+        if isinstance(obj, xr.DataArray) and isinstance(obj.data, cubed.Array):
+            computed_data = cubed.compute(obj.data)[0]
+            return xr.DataArray(
+                computed_data,
+                coords=obj.coords,
+                dims=obj.dims,
+                name=obj.name,
+                attrs=obj.attrs,
+            )
+        elif isinstance(obj, xr.Dataset):
+            computed_vars = {}
+            for k, v in obj.data_vars.items():
+                if isinstance(v.data, cubed.Array):
+                    computed_vars[k] = _compute_lazy_aware(v)
+                else:
+                    computed_vars[k] = v
+            return xr.Dataset(
+                data_vars=computed_vars, coords=obj.coords, attrs=obj.attrs
+            )
+        elif isinstance(obj, cubed.Array):
+            return cubed.compute(obj)[0]
+
+    if is_dask(obj):
+        if dask is None:
+            return obj
+        return dask.compute(obj)[0]
+
+    return obj
+
+
 def _get_array_namespace(*objs: Any) -> Any:
     """
     Get the appropriate array namespace (numpy, dask.array, or cubed) for the given objects.
@@ -1327,7 +1409,7 @@ def create_grid_like(
                     tasks_dict["ymin"] = y_min
                     tasks_dict["ymax"] = y_max
 
-                results = dask.compute(tasks_dict)[0]
+                results = _compute_lazy_aware(tasks_dict)
                 extent = (
                     float(results.get("xmin", x_min)),
                     float(results.get("xmax", x_max)),
@@ -1361,7 +1443,7 @@ def create_grid_like(
                 if y_da.size > 1:
                     tasks_dict["res_y"] = abs(y_da.diff(y_da.dims[0]).mean())
 
-                results = dask.compute(tasks_dict)[0]
+                results = _compute_lazy_aware(tasks_dict)
                 x_min_val = float(results.get("x_min", x_min))
                 x_max_val = float(results.get("x_max", x_max))
                 y_min_val = float(results.get("y_min", y_min))
@@ -1436,7 +1518,7 @@ def create_grid_like(
                     tasks_dict["lon_min"] = lon_min
                     tasks_dict["lon_max"] = lon_max
 
-                results = dask.compute(tasks_dict)[0]
+                results = _compute_lazy_aware(tasks_dict)
                 lat_range = (
                     float(results.get("lat_min", lat_min)),
                     float(results.get("lat_max", lat_max)),
@@ -1472,7 +1554,7 @@ def create_grid_like(
                 if lon_da.size > 1:
                     tasks_dict["res_lon"] = abs(lon_da.diff(lon_da.dims[-1]).mean())
 
-                results = dask.compute(tasks_dict)[0]
+                results = _compute_lazy_aware(tasks_dict)
                 lat_min_val = float(results.get("lat_min", lat_min))
                 lat_max_val = float(results.get("lat_max", lat_max))
                 lon_min_val = float(results.get("lon_min", lon_min))

tests/test_aero_compute_agnostic.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+import numpy as np
+import xarray as xr
+
+from xregrid.utils import _compute_lazy_aware, is_cubed, is_dask, is_lazy
+
+
+def test_is_dask_and_is_cubed_nested() -> None:
+    """
+    Test recursive detection of dask and cubed arrays in nested structures.
+    """
+    # 1. Eager arrays
+    da_eager = xr.DataArray(np.arange(10))
+    nested_eager = {"a": [da_eager], "b": (da_eager,)}
+
+    assert not is_dask(da_eager)
+    assert not is_dask(nested_eager)
+    assert not is_cubed(da_eager)
+    assert not is_cubed(nested_eager)
+    assert not is_lazy(nested_eager)
+
+    # 2. Lazy (Dask) arrays
+    da_lazy = da_eager.chunk(5)
+    nested_lazy = {"a": [da_lazy], "b": (da_lazy,)}
+
+    assert is_dask(da_lazy)
+    assert is_dask(nested_lazy)
+    assert is_lazy(nested_lazy)
+
+    # Cubed is not installed or active, but shouldn't raise errors
+    assert not is_cubed(da_lazy)
+    assert not is_cubed(nested_lazy)
+
+
+def test_compute_lazy_aware_eager_and_lazy() -> None:
+    """
+    Test the _compute_lazy_aware utility for eager and lazy data.
+
+    Verifies the logic twice:
+    - Once with eager data (returns immediately without changing structure)
+    - Once with lazy (Dask) data (fully computes and returns eager data)
+    """
+    # 1. Eager NumPy Path
+    eager_val = xr.DataArray(np.arange(5))
+    eager_nested = {"a": eager_val, "b": [eager_val, 42], "c": (eager_val,)}
+
+    res_eager = _compute_lazy_aware(eager_nested)
+    assert not is_lazy(res_eager)
+    assert isinstance(res_eager["a"], xr.DataArray)
+    np.testing.assert_array_equal(res_eager["a"].values, np.arange(5))
+    np.testing.assert_array_equal(res_eager["b"][0].values, np.arange(5))
+    assert res_eager["b"][1] == 42
+    np.testing.assert_array_equal(res_eager["c"][0].values, np.arange(5))
+
+    # 2. Lazy Dask Path
+    lazy_val = eager_val.chunk(2)
+    lazy_nested = {"a": lazy_val, "b": [lazy_val, 42], "c": (lazy_val,)}
+
+    assert is_lazy(lazy_nested)
+
+    res_lazy = _compute_lazy_aware(lazy_nested)
+    assert not is_lazy(res_lazy)
+    assert isinstance(res_lazy["a"], xr.DataArray)
+    np.testing.assert_array_equal(res_lazy["a"].values, np.arange(5))
+    np.testing.assert_array_equal(res_lazy["b"][0].values, np.arange(5))
+    assert res_lazy["b"][1] == 42
+    np.testing.assert_array_equal(res_lazy["c"][0].values, np.arange(5))