Skip to content

calling xarray data type methods

@finn-heukamp

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.