Fix and extend wrapper

windeval/tests/test_wrapper.py

-import windeval
+from windeval import wrapper
 
 
 def test_ekman(X):
-    windeval.ekman(X, drag_coefficient="large_and_pond_1981", extend_ranges=True)
+    wrapper.ekman(X, drag_coefficient="large_and_pond_1981", extend_ranges=True)
 
 
 def test_sverdrup(X):
-    windeval.sverdrup(X, drag_coefficient="large_and_yeager_2004")
+    wrapper.sverdrup(X, drag_coefficient="large_and_yeager_2004")

windeval/wrapper.py

 """Wrapper."""
 
 import xarray as xr
-from typing import Optional
-from . import processing as pr
+from functools import singledispatch
+from typing import Optional, Any, Dict, List
+from . import processing
+
+
+@singledispatch
+def calculate(ds, *args, **kwargs):
+    raise NotImplementedError("Data type not supported.")
+
+
+@calculate.register
+def _(
+    ds: xr.Dataset, var: str, diag: str, *args: Any, **kwargs: Dict[str, Any]
+) -> xr.DataArray:
+    return calculate(ds[var], diag, *args, **kwargs)
+
+
+@calculate.register  # type: ignore
+def _(
+    da: xr.DataArray, diag: str, *args: Any, **kwargs: Dict[str, Any]
+) -> xr.DataArray:
+
+    f = getattr(processing, diag, None)
+    if f is not None:
+        y = f(da, *args, **kwargs)
+    else:
+        y = getattr(da, diag)(*args, **kwargs)
+
+    return y
+
+
+@calculate.register  # type: ignore
+def _(
+    d: dict, keys: List[str], diag: str, *args: Any, **kwargs: Dict[str, Any]
+) -> xr.DataArray:
+
+    f = getattr(processing, diag)
+    return f(*[d[k] for k in keys], *args, **kwargs)
 
 
 def ekman(
@@ -40,8 +76,8 @@ def ekman(
             continue
         d[x] = locals()[x]
     c = "northward" if component == "eastward" else "eastward"
-    getattr(pr, "surface_downward_" + c + "_stress")(X, **d)
-    getattr(pr, component + "_ekman_transport")(X)
+    getattr(processing, "surface_downward_" + c + "_stress")(X, **d)
+    getattr(processing, component + "_ekman_transport")(X)
 
     return X
 
@@ -77,8 +113,8 @@ def sverdrup(
         if locals()[x] is None:
             continue
         d[x] = locals()[x]
-    pr.surface_downward_eastward_stress(X, **d)
-    pr.surface_downward_northward_stress(X, **d)
-    pr.sverdrup_transport(X)
+    processing.surface_downward_eastward_stress(X, **d)
+    processing.surface_downward_northward_stress(X, **d)
+    processing.sverdrup_transport(X)
 
     return X