Create a decorator to check units of function arguments.

def declare_units(out_units, **units_by_name):
    """Create a decorator to check units of function arguments."""

    def dec(func):
        # Match the signature of the function to the arguments given to the decorator
        sig = signature(func)
        bound_units = sig.bind_partial(**units_by_name)

        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            # Match all passed in value to their proper arguments so we can check units
            bound_args = sig.bind(*args, **kwargs)
            for name, val in bound_args.arguments.items():
                _check_units(val, bound_units.arguments.get(name, None))

            out = func(*args, **kwargs)

            # In the generic case, we use the default units that should have been propagated by the computation.
            if '[' in out_units:
                _check_units(out, out_units)

            # Otherwise, we specify explicitly the units.
            else:
                out.attrs['units'] = out_units
            return out

        return wrapper

    return dec