calling xarray data type methods
You wrote several mappings for xarray data type methods, you could combine all of them into functional calls, e.g. like this
calculate(ds, "wind_speed", "std", dim="time")
# or
calculate(ds["wind_speed"], "mean", "time")
with the below code
import xarray as xr
from functools import singledispatch
from typing import Optional, Any, Dict
@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 getattr(ds[var], diag)(*args, **kwargs)
@calculate.register # type: ignore
def _(
da: xr.DataArray, diag: str, *args: Any, **kwargs: Dict[str, Any]
) -> xr.DataArray:
return getattr(da, diag)(*args, **kwargs)
I've not yet added this to the repo, however, if you want something like this to exist, we could add it to the wrapper module I've added lately.