def check(ureg, *args):
"""Decorator to for quantity type checking for function inputs.
Use it to ensure that the decorated function input parameters match
the expected type of pint quantity.
Use None to skip argument checking.
:param ureg: a UnitRegistry instance.
:param args: iterable of input units.
:return: the wrapped function.
:raises:
:class:`DimensionalityError` if the parameters don't match dimensions
"""
dimensions = [ureg.get_dimensionality(dim) if dim is not None else None for dim in args]
def decorator(func):
assigned = tuple(attr for attr in functools.WRAPPER_ASSIGNMENTS if hasattr(func, attr))
updated = tuple(attr for attr in functools.WRAPPER_UPDATES if hasattr(func, attr))
@functools.wraps(func, assigned=assigned, updated=updated)
def wrapper(*values, **kwargs):
for dim, value in zip_longest(dimensions, values):
if dim is None:
continue
val_dim = ureg.get_dimensionality(value)
if val_dim != dim:
raise DimensionalityError(value, 'a quantity of',
val_dim, dim)
return func(*values, **kwargs)
return wrapper
return decorator
评论列表
文章目录