"""
This module is a place to put decorators used by methods in aurora.
It is a work in progress.
Development notes:
| Here is the decorator pattern
| def decorator(func):
| @functools.wraps(func)
| def wrapper_decorator(\*args, \*\*kwargs):
| # Do something before
| value = func\*args, \*\*kwargs)
| # Do something after
| return value
| return wrapper_decorator
"""
import functools
import xarray as xr
[docs]def can_use_xr_dataarray(func):
"""
Decorator to allow functions that operate on xr.Dataset to also operate
on xr.DataArray, by converting xr.DataArray to xr.Dataset, calling the
function, and then converting back to xr.DataArray
Most of the windowed time series methods are
intended to work with xarray.Dataset class. But I would like to be able
to pass them xarray.DataArray objects. This class casts a DataArray to a
Dataset, runs it through func and casts back to a DataArray.
A similar decorator could be written for numpy arrays.
Parameters
----------
func
Returns
-------
"""
@functools.wraps(func)
def wrapper_decorator(*args, **kwargs):
if isinstance(kwargs["data"], xr.DataArray):
kwargs["data"] = kwargs["data"].to_dataset("channel")
input_was_dataarray = True
else:
input_was_dataarray = False
if isinstance(func, staticmethod):
callable_func = func.__get__(None, object)
processed_obj = callable_func(*args, **kwargs)
else:
processed_obj = func(*args, **kwargs)
if input_was_dataarray:
processed_obj = processed_obj.to_array("channel")
return processed_obj
return wrapper_decorator