Skip to main content
`Ctrl`+`K`

# Wrapping and checking functions

In some cases you might want to use pint with a pre-existing web service or library which is not units aware. Or you might want to write a fast implementation of a numerical algorithm that requires the input values in some specific units.

For example, consider a function to return the period of the pendulum within a hypothetical physics library. The library does not use units, but instead requires you to provide numerical values in certain units:

    >>> from simple_physics import pendulum_period
    >>> help(pendulum_period)
    Help on function pendulum_period in module simple_physics:

    pendulum_period(length)
    Return the pendulum period in seconds. The length of the pendulum
    must be provided in meters.

    >>> pendulum_period(1)
    2.0064092925890407

This behaviour is very error prone, in particular when combining multiple libraries. You could wrap this function to use Quantities instead:

    >>> from pint import UnitRegistry
    >>> ureg = UnitRegistry()
    >>> Q_ = ureg.Quantity
    >>> def mypp_caveman(length):
    ...     return pendulum_period(length.to(ureg.meter).magnitude) * ureg.second

and:

    >>> mypp_caveman(100 * ureg.centimeter)
    <Quantity(2.00640929, 'second')>

Pint provides a more convenient way to do this:

    >>> mypp = ureg.wraps(ureg.second, ureg.meter)(pendulum_period)

Or in the decorator format:

    >>> @ureg.wraps(ureg.second, ureg.meter)
    ... def mypp(length):
    ...     return pendulum_period(length)
    >>> mypp(100 * ureg.centimeter)
    <Quantity(2.00640929, 'second')>

wraps takes 3 input arguments:

> - ret: the return units.

>     Use None to skip conversion.

> - args: the inputs units for each argument, as an iterable.

>     Use None to skip conversion of any given element.

> - strict: if True all convertible arguments must be a Quantity

>     and others will raise a ValueError (True by default)

## Strict Mode

By default, the function is wrapped in strict mode. In this mode, the input arguments assigned to units must be a Quantities.

    >>> mypp(1. * ureg.meter)
    <Quantity(2.00640929, 'second')>
    >>> mypp(1.)
    Traceback (most recent call last):
    ...
    ValueError: A wrapped function using strict=True requires quantity for all arguments with not None units. (error found for meter, 1.0)

To enable using non-Quantity numerical values, set strict to False`.

    >>> mypp_ns = ureg.wraps(ureg.second, ureg.meter, False)(pendulum_period)
    >>> mypp_ns(1. * ureg.meter)
    <Quantity(2.00640929, 'second')>
    >>> mypp_ns(1.)
    <Quantity(2.00640929, 'second')>

In this mode, the value is assumed to have the correct units.

## Multiple arguments or return values

For a function with more arguments, use a tuple:

    >>> from simple_physics import pendulum_period2
    >>> help(pendulum_period2)
    Help on function pendulum_period2 in module simple_physics:

    pendulum_period2(length, swing_amplitude)
    Return the pendulum period in seconds. The length of the pendulum
    must be provided in meters. The swing_amplitude must be in radians.

    >>> mypp2 = ureg.wraps(ureg.second, (ureg.meter, ureg.radians))(pendulum_period2)
    ...

Or if the function has multiple outputs:

    >>> mypp3 = ureg.wraps((ureg.second, ureg.meter / ureg.second),
    ...                    (ureg.meter, ureg.radians))(pendulum_period_maxspeed)
    ...

If there are more return values than specified units, `None` is assumed for the extra outputs. For example, given the NREL SOLPOS calculator that outputs solar zenith, azimuth and air mass, the following wrapper assumes no units for airmass

    @ureg.wraps(('deg', 'deg'), ('deg', 'deg', 'millibar', 'degC'))
    def solar_position(lat, lon, press, tamb, timestamp):
        return zenith, azimuth, airmass

## Optional arguments

For a function with named keywords with optional values, use a tuple for all arguments:

    >>> @ureg.wraps(ureg.second, (ureg.meters, ureg.meters/ureg.second**2, None))
    ... def calculate_time_to_fall(height, gravity=Q_(9.8, 'm/s^2'), verbose=False):
    ...     """Calculate time to fall from a height h.
    ...
    ...     By default, the gravity is assumed to be earth gravity,
    ...     but it can be modified.
    ...
    ...     d = .5 * g * t**2
    ...     t = sqrt(2 * d / g)
    ...     """
    ...     t = math.sqrt(2 * height / gravity)
    ...     if verbose: print(str(t) + " seconds to fall")
    ...     return t
    ...
    >>> lunar_module_height = Q_(22, 'feet') + Q_(11, 'inches')
    >>> calculate_time_to_fall(lunar_module_height, verbose=True)
    1.1939473204801092 seconds to fall
    <Quantity(1.19394732, 'second')>
    >>> moon_gravity = Q_(1.625, 'm/s^2')
    >>> calculate_time_to_fall(lunar_module_height, gravity=moon_gravity)
    <Quantity(2.932051, 'second')>

## Specifying relations between arguments

In certain cases, you may not be concerned with the actual units and only care about the unit relations among arguments.

This is done using a string starting with the equal sign =:

    >>> @ureg.wraps('=A**2', ('=A', '=A'))
    ... def sqsum(x, y):
    ...     return x * x  + 2 * x * y + y * y

which can be read as the first argument (x) has certain units (we labeled them A), the second argument (y) has the same units as the first (A again). The return value has the unit of x squared (A**2)

You can use more than one label:

    >>> @ureg.wraps('=A**2*B', ('=A', '=A*B', '=B'))
    ... def some_function(x, y, z):
    ...     pass

With optional arguments

    >>> @ureg.wraps('=A*B', ('=A', '=B'))
    ... def get_displacement(time, rate=Q_(1, 'm/s')):
    ...     return time * rate
    ...
    >>> get_displacement(Q_(2, 's'))
    <Quantity(2, 'meter')>
    >>> get_displacement(Q_(2, 's'), Q_(1, 'deg/s'))
    <Quantity(2, 'degree')>

## Ignoring an argument or return value

To avoid the conversion of an argument or return value, use None

    >>> mypp3 = ureg.wraps((ureg.second, None), ureg.meter)(pendulum_period_error)

## Checking dimensionality

When you want pint quantities to be used as inputs to your functions, pint provides a wrapper to ensure units are of correct type - or more precisely, they match the expected dimensionality of the physical quantity.

Similar to wraps(), you can pass None to skip checking of some parameters, but the return parameter type is not checked.

    >>> mypp = ureg.check('[length]')(pendulum_period)

In the decorator format:

    >>> @ureg.check('[length]')
    ... def pendulum_period(length):
    ...     return 2*math.pi*math.sqrt(length/G)

If you just want to check the dimensionality of a quantity, you can do so with the built-in ‘check’ function.

    >>> distance = 1 * ureg.m
    >>> distance.check('[length]')
    True
    >>> distance.check('[time]')
    False