Force pint args to be of particular dimensionality

pint_args_decorator.py

import pint
import functools
import itertools

ureg = pint.UnitRegistry()


def require_args(*args):
    units = args

    def decorator(func):
        @functools.wraps(func)
        def wrapper(*values, **kwargs):
            for dim, value in itertools.izip_longest(units, values):
                if value.dimensionality != ureg.get_dimensionality(dim):
                    raise TypeError(
                        'Expected units of %s, got %s' %
                        (dim, value.dimensionality))
            result = func(*values, **kwargs)
            return result
        return wrapper
    return decorator


@require_args('[speed]', '[length]')
def time_taken(speed, distance):
    return distance / speed


@require_args('[acceleration]', '[time]')
def get_kessel_run(acceleration, time):
    return 0.5 * acceleration * time ** 2

if __name__ == "__main__":
    from pint.compat import unittest
    from pint.testsuite import QuantityTestCase

    class TestRequireArgs(QuantityTestCase):
        def test_time(self):
            elapsed = time_taken(40.0 * ureg.mm / ureg.minute, ureg.meter)
            self.assertAlmostEqual(elapsed, 1500 * ureg.second)

        def test_space(self):
            wookie_drive = 3968210.19805 * ureg.meter / ureg.second ** 2
            ewoks = get_kessel_run(wookie_drive, 5 * ureg.day).to('parsec')
            self.assertQuantityAlmostEqual(ewoks, 12.0 * ureg.parsec)

        def test_no_return(self):
            with self.assertRaises(TypeError):
                get_kessel_run(ureg.kilogram * 1, ureg.angstrom * 4)

    unittest.main()