expects decorator

pint_xarray/__init__.py

@@ -8,6 +8,7 @@
 from . import accessors, formatting, testing  # noqa: F401
 from .accessors import default_registry as unit_registry
 from .accessors import setup_registry
+from .checking import expects  # noqa: F401
 
 try:
     __version__ = version("pint-xarray")

pint_xarray/checking.py

@@ -0,0 +1,118 @@
+import functools
+
+from pint import Quantity
+from xarray import DataArray
+
+from .accessors import PintDataArrayAccessor
+
+
+def expects(*args_units, return_units=None, **kwargs_units):
+    """
+    Decorator which checks the inputs and outputs of the decorated function have certain units.
+
+    Arguments
+
+    Note that the coordinates of input DataArrays are not checked, only the data.
+    So if your decorated function uses coordinates and you wish to check their units,
+    you should pass the coordinates of interest as separate arguments.
+
+    Parameters
+    ----------
+    func: function
+        Decorated function, which accepts zero or more xarray.DataArrays or pint.Quantitys as inputs,
+        and may optionally return one or more xarray.DataArrays or pint.Quantitys.
+    args_units : Union[str, pint.Unit, None]
+        Unit to expect for each positional argument given to func.
+
+        A value of None indicates not to check that argument for units (suitable for flags
+        and other non-data arguments).
+    return_units : Union[Union[str, pint.Unit, None, False], Sequence[Union[str, pint.Unit, None]], Optional
+        The expected units of the returned value(s), either as a single unit or as an iterable of units.
+
+        A value of None indicates not to check that return value for units (suitable for flags and other
+        non-data arguments). Passing False means that no return value is expected from the function at all,
+        and an error will be raised if a return value is found.
+    kwargs_units : Dict[str, Union[str, pint.Unit, None]], Optional
+        Unit to expect for each keyword argument given to func.
+
+        A value of None indicates not to check that argument for units (suitable for flags
+        and other non-data arguments).
+
+    Returns
+    -------
+    return_values
+        Return values of the wrapped function, either a single value or a tuple of values.
+
+    Raises
+    ------
+    TypeError
+        If an argument or return value has a specified unit, but is not an xarray.DataArray.
+
+
+    Examples
+    --------
+
+    Decorating a function which takes one quantified input, but returns a non-data value (in this case a boolean).
+
+    >>> @expects('deg C')
+    ... def above_freezing(temp):
+    ...     return temp > 0
+    ...
+
+
+    TODO: example where we check units of an optional weighted kwarg
+    """
+
+    # TODO: Check args_units, kwargs_units, and return_units types
+    # TODO: Check number of arguments line up
+
+    def _expects_decorator(func):
+
+        @functools.wraps(func)
+        def _unit_checking_wrapper(*args, **kwargs):
+
+            converted_args = []
+            for arg, arg_unit in zip(args, args_units):
+                converted_arg = _check_then_convert_to(arg, arg_unit)
+                converted_args.append(converted_arg)
+
+            converted_kwargs = {}
+            for key, val in kwargs.items():
+                kwarg_unit = kwargs_units[key]
+                converted_kwargs[key] = _check_then_convert_to(val, kwarg_unit)
+
+            results = func(*converted_args, **converted_kwargs)
+
+            if results is not None:
+                if return_units is False:
+                    raise ValueError("Did not expect function to return anything")
+                elif return_units is not None:
+                    # TODO check something was actually returned
+                    # TODO check same number of things were returned as expected
+                    # TODO handle single return value vs tuple of return values
+
+                    converted_results = []
+                    for return_unit, return_value in zip(return_units, results):
+                        converted_result = _check_then_convert_to(return_value, return_unit)
+                        converted_results.append(converted_result)
+
+                    return tuple(converted_results)
+                else:
+                    return results
+            else:
+                if return_units:
+                    raise ValueError("Expected function to return something")
+
+        return _unit_checking_wrapper
+
+    return _expects_decorator
+
+
+def _check_then_convert_to(obj, units):
+    if isinstance(obj, Quantity):
+        return obj.to(units)
+    elif isinstance(obj, DataArray):
+        return obj.pint.to(units)
+    else:
+        raise TypeError("Can only expect units for arguments of type xarray.DataArray or pint.Quantity,"
+                        f"not {type(obj)}")

pint_xarray/tests/test_checking.py

@@ -0,0 +1,97 @@
+import pytest
+import pint
+import xarray as xr
+
+from pint import UnitRegistry
+
+from ..checking import expects
+
+ureg = UnitRegistry()
+
+
+class TestExpects:
+    def test_single_arg(self):
+
+        @expects('degC')
+        def above_freezing(temp : pint.Quantity):
+            return temp.magnitude > 0
+
+        f_q = pint.Quantity(20, units='degF')
+        assert above_freezing(f_q) == False
+
+        c_q = pint.Quantity(-2, units='degC')
+        assert above_freezing(c_q) == False
+
+        @expects('degC')
+        def above_freezing(temp : xr.DataArray):
+            return temp.pint.magnitude > 0
+
+        f_da = xr.DataArray(20).pint.quantify(units='degF')
+        assert above_freezing(f_da) == False
+
+        c_da = xr.DataArray(-2).pint.quantify(units='degC')
+        assert above_freezing(c_da) == False
+
+    def test_single_kwarg(self):
+
+        @expects('meters', c='meters / second')
+        def freq(wavelength, c=None):
+            if c is None:
+                c = ureg.speed_of_light
+
+            return c / wavelength
+
+    def test_single_return_value(self):
+
+        @expects('Hz')
+        def period(freq):
+            return 1 / freq
+
+        f = pint.Quantity(10, units='Hz')
+
+        # test conversion
+        T = period(f)
+        assert f.units == 'seconds'
+
+        # test wrong dimensions for conversion
+        ...
+
+    @pytest.mark.xfail
+    def test_multiple_return_values(self):
+        raise NotImplementedError
+
+    @pytest.mark.xfail
+    def test_mixed_args_kwargs_return_values(self):
+        raise NotImplementedError
+
+    @pytest.mark.xfail
+    def test_invalid_input_types(self):
+        raise NotImplementedError
+
+    @pytest.mark.xfail
+    def test_invalid_return_types(self):
+        raise NotImplementedError
+
+    @pytest.mark.xfail
+    def test_unquantified_arrays(self):
+        raise NotImplementedError
+
+    @pytest.mark.xfail
+    def test_wrong_number_of_args(self):
+        raise NotImplementedError
+
+    @pytest.mark.xfail
+    def test_nonexistent_kwarg(self):
+        raise NotImplementedError
+
+    @pytest.mark.xfail
+    def test_unexpected_return_value(self):
+        raise NotImplementedError
+
+    @pytest.mark.xfail
+    def test_expected_return_value(self):
+        raise NotImplementedError
+
+    @pytest.mark.xfail
+    def test_wrong_number_of_return_values(self):
+        raise NotImplementedError