check_units

plasmapy.utils.decorators.checks.check_units(func=None, checks_on_return: dict[str, Any] | None = None, **checks: dict[str, Any])[source]

A decorator to ‘check’ — limit/control — the units of input and return arguments to a function or method.

Parameters:

func – The function to be decorated
checks_on_return (list of units or dict of unit specifications) – Specifications for unit checks on the return of the function being wrapped. (see “check units”_ for valid specifications)

**checks (list of units or dict of unit specifications) –

Specifications for unit checks on the input arguments of the function being wrapped. Each keyword argument in checks is the name of a function argument to be checked and the keyword value contains the unit check specifications.

Unit checks can be defined by passing one of the astropy units, a list of astropy units, or a dictionary containing the keys defined below. Units can also be defined with function annotations, but must be consistent with decorator **checks arguments if used concurrently. If a key is omitted, then the default value will be assumed.

Key	Type	Description
units		list of desired astropy `units`
equivalencies		[DEFAULT `None`] A list of equivalent pairs to try if the units are not directly convertible. (see `equivalencies`, and/or astropy equivalencies)
pass_equivalent_units	`bool`	[DEFAULT `False`] allow equivalent units to pass

Notes

Checking of function arguments *args and **kwargs is not supported.
Decorator does NOT perform any unit conversions, look to validate_quantities() if that functionality is desired.
If it is desired that None values do not raise errors or warnings, then include None in the list of units or as a default value for the function argument.
If units are not specified in checks, then the decorator will attempt to identify desired units by examining the function annotations.
Full functionality is defined by the class CheckUnits.

Examples

Define units with decorator parameters:

import astropy.units as u
from plasmapy.utils.decorators import check_units

@check_units(arg1={"units": u.cm}, arg2=u.cm, checks_on_return=[u.cm, u.km])
def foo(arg1, arg2):
    return arg1 + arg2


# or on a method
class Foo:
    @check_units(arg1={"units": u.cm}, arg2=u.cm, checks_on_return=[u.cm, u.km])
    def bar(self, arg1, arg2):
        return arg1 + arg2

Define units with function annotations:

@check_units
def foo(arg1: u.cm, arg2: u.cm) -> u.cm:
    return arg1 + arg2


# or on a method
class Foo:
    @check_units
    def bar(self, arg1: u.cm, arg2: u.cm) -> u.cm:
        return arg1 + arg2

Allow None values to pass:

@check_units(checks_on_return=[u.cm, None])
def foo(arg1: u.cm = None):
    return arg1

Allow return values to have equivalent units:

@check_units(
    arg1={"units": u.cm},
    checks_on_return={"units": u.km, "pass_equivalent_units": True},
)
def foo(arg1):
    return arg1

Allow equivalent units to pass with specified equivalencies:

@check_units(
    arg1={
        "units": u.K,
        "equivalencies": u.temperature(),
        "pass_equivalent_units": True,
    }
)
def foo(arg1):
    return arg1