ranged_quantity_input

pycraf.utils.ranged_quantity_input(func=None, **kwargs)

A decorator for validating the units of arguments to functions.

This decorator was adapted from Astropy’s quantity_input, but adds range checking and the possibilities to strip the units before feeding into the decorated function. It also allows to apply a new unit to the returned value (quantity_input only does this in conjuction with type annotations).

A UnitsError will be raised if the unit attribute of the argument is not equivalent to the unit specified to the decorator or in the annotation. If the argument has no unit attribute, i.e. it is not a Quantity object, a ValueError will be raised.

Where an equivalency is specified in the decorator, the function will be executed with that equivalency in force.

Parameters:

funcfunction

The function to decorate.

**kwargsany number of key word arguments

The function argument names and ranges that are to be checked for the decorated function. Must have the form param=(min, max, unit), e.g.:

@ranged_quantity_input(a=(0, 1, u.m), b=(0, None, u.s))
def func(a, b):
    return a ** 2, 1 / b

will check that input a has unit of meters (or equivalent) and is in the range between zero and one meters; and that b is at least zero seconds.

equivalencieslist of functions

Equivalencies functions to apply (see Astropy docs).

strip_input_unitsbool, optional

Whether to strip units from parameters. Only applied to parameters that are “registered” in the decorator, see examples. (default: False)

output_unitUnit or tuple of Unit, optional

Add units to the return value(s) of the decorated function. Note that internally the given units are multiplied with the return values, which means you should only use this if you have stripped the units from the input (or otherwise made sure that the return values are unit-less).

allow_nonebool, optional

Allow to use None as default value; see examples.

Returns:

ranged_quantity_inputfunction decorator

Function decorator to check units and value ranges.

Notes

The checking of arguments inside variable arguments to a function is not supported (i.e. *arg or **kwargs).

Examples

In the most basic form, ranged_quantity_input behaves like quantity_input, but adds range checking:

>>> from pycraf.utils import ranged_quantity_input
>>> import astropy.units as u

>>> @ranged_quantity_input(a=(0, 1, u.m))
... def func(a):
...     return a ** 2

>>> func(0.5 * u.m)  
<Quantity 0.25 m2>

>>> func(2 * u.m)
Traceback (most recent call last):
...
ValueError: Argument 'a' to function 'func' out of range
(allowed 0 to 1 m).

It is possible to disable range checking, for the lower, upper, or both bounds, e.g.:

>>> @ranged_quantity_input(a=(0, None, u.m))
... def func(a):
...     return a ** 2

>>> func(2 * u.m)  
<Quantity 4. m2>

Often one wants to add units support to third-party functions, which expect simple types:

>>> # this is defined somewhere else
>>> def _func(a):
...     assert isinstance(a, float), 'No Way!'
...     return a ** 2

>>> _func(0.5 * u.m)
Traceback (most recent call last):
...
AssertionError: No Way!

We can do the following to the rescue:

>>> @ranged_quantity_input(a=(0, 1, u.m), strip_input_units=True)
... def func(a):
...     return _func(a)

>>> # which is the same as
>>> # func = ranged_quantity_input(
>>> #    a=(0, 1, u.m), strip_input_units=True
>>> #    )(_func)

>>> func(0.5 * u.m)  
np.float64(0.25)

However, by doing this there are still no units for the output. We can fix this with the output_unit option:

>>> @ranged_quantity_input(
...     a=(0, 1, u.m),
...     strip_input_units=True,
...     output_unit=u.m ** 2
...     )
... def func(a):
...     return _func(a)

>>> func(0.5 * u.m)  
<Quantity 0.25 m2>

If you have several return values (tuple), just provide a tuple of output units.

The decorator also works flawlessly with default values:

>>> @ranged_quantity_input(a=(0, 1, u.m))
... def func(a=0.5 * u.m):
...     return a ** 2

>>> func()  
<Quantity 0.25 m2>

However, sometimes one wants to use None as default, which will fail, because None has no unit:

>>> @ranged_quantity_input(a=(0, 1, u.m))
... def func(a=None):
...     return a ** 2

>>> func()
Traceback (most recent call last):
...
TypeError: Argument 'a' to function 'func' has no 'unit'
attribute. You may want to pass in an astropy Quantity instead.

One can use the allow_none option, to deal with such cases:

>>> @ranged_quantity_input(a=(0, 1, u.m), allow_none=True)
... def func(a=None):
...     if a is None:
...         a = 0.5
...     return a ** 2

>>> func()  
0.25

and of course, the unit check still works, if a something other than None is provided:

>>> func(1 * u.s)  
Traceback (most recent call last):
...
UnitsError: Argument 'a' to function
'func' must be in units convertible to 'm'.