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.

Assignee Loading
Time tracking Loading