"""
pint.registry_helpers
~~~~~~~~~~~~~~~~~~~~~
Miscellaneous methods of the registry written as separate functions.
:copyright: 2016 by Pint Authors, see AUTHORS for more details..
:license: BSD, see LICENSE for more details.
"""
from __future__ import annotations
import functools
from collections.abc import Callable, Iterable
from inspect import Parameter, signature
from itertools import zip_longest
from typing import TYPE_CHECKING, Any, overload
from ._typing import FuncType
from .errors import DimensionalityError
from .util import UnitsContainer, to_units_container
if TYPE_CHECKING:
from ._typing import Quantity, Unit
from .facets.plain import GenericPlainRegistry
from .registry import UnitRegistry
def _replace_units(original_units, values_by_name):
"""Convert a unit compatible type to a UnitsContainer.
Parameters
----------
original_units :
a UnitsContainer instance.
values_by_name :
a map between original names and the new values.
Returns
-------
"""
q = UnitsContainer({})
for arg_name, exponent in original_units.items():
q *= getattr(values_by_name[arg_name], "_units", UnitsContainer({})) ** exponent
return q
def _to_units_container(a, registry=None):
"""Convert a unit compatible type to a UnitsContainer,
checking if it is string field prefixed with an equal
(which is considered a reference)
Parameters
----------
a :
registry :
(Default value = None)
Returns
-------
UnitsContainer, bool
"""
if isinstance(a, str) and "=" in a:
return to_units_container(a.split("=", 1)[1]), True
return to_units_container(a, registry), False
def _parse_wrap_args(args, registry=None, *, has_varargs=False):
# Arguments which contain definitions
# (i.e. names that appear alone and for the first time)
defs_args = set()
defs_args_ndx = set()
# Arguments which depend on others
dependent_args_ndx = set()
# Arguments which have units.
unit_args_ndx = set()
# _to_units_container
args_as_uc = [_to_units_container(arg, registry) for arg in args]
# Check for references in args, remove None values
for ndx, (arg, is_ref) in enumerate(args_as_uc):
if arg is None:
continue
elif is_ref:
if len(arg) == 1:
[(key, value)] = arg.items()
if value == 1 and key not in defs_args:
# This is the first time that
# a variable is used => it is a definition.
defs_args.add(key)
defs_args_ndx.add(ndx)
args_as_uc[ndx] = (key, True)
else:
# The variable was already found elsewhere,
# we consider it a dependent variable.
dependent_args_ndx.add(ndx)
else:
dependent_args_ndx.add(ndx)
else:
unit_args_ndx.add(ndx)
# Check that all valid dependent variables
for ndx in dependent_args_ndx:
arg, is_ref = args_as_uc[ndx]
if not isinstance(arg, dict):
continue
if not set(arg.keys()) <= defs_args:
raise ValueError(
"Found a missing token while wrapping a function: "
"Not all variable referenced in %s are defined using !" % args[ndx]
)
single_indices = tuple((ndx,) for ndx in range(len(args)))
def _converter(ureg, sig, values, kw, strict):
len_initial_values = len(values)
if has_varargs:
# Bind before flattening: keyword-only arguments must not be
# mistaken for elements of *args, even when defaults are omitted.
bound = sig.bind(*values, **kw)
bound.apply_defaults()
values = []
indices = []
for name, param in sig.parameters.items():
if param.kind == Parameter.VAR_KEYWORD:
continue
start = len(values)
if param.kind == Parameter.VAR_POSITIONAL:
values.extend(bound.arguments[name])
else:
values.append(bound.arguments[name])
indices.append(range(start, len(values)))
else:
indices = single_indices
# pack kwargs
for i, (param_name, param) in enumerate(sig.parameters.items()):
if param.kind == Parameter.VAR_KEYWORD:
continue
if i >= len_initial_values:
values.append(kw[param_name])
values_by_name = {}
# first pass: Grab named values
for ndx in defs_args_ndx:
positions = indices[ndx]
# An empty *args has no unit-bearing value to bind. Its reference
# is dimensionless, just as for a plain scalar argument.
value = values[positions[0]] if positions else None
values_by_name[args_as_uc[ndx][0]] = value
if positions:
values[positions[0]] = getattr(value, "_magnitude", value)
# Only the first variadic value defines the reference; every
# subsequent one must be converted to that same unit.
for pos in positions[1:]:
other = values[pos]
values[pos] = ureg._convert(
getattr(other, "_magnitude", other),
getattr(other, "_units", UnitsContainer({})),
getattr(value, "_units", UnitsContainer({})),
)
# second pass: calculate derived values based on named values
for ndx in dependent_args_ndx:
positions = indices[ndx]
if not positions:
continue
target_units = _replace_units(args_as_uc[ndx][0], values_by_name)
assert target_units is not None
for pos in positions:
value = values[pos]
values[pos] = ureg._convert(
getattr(value, "_magnitude", value),
getattr(value, "_units", UnitsContainer({})),
target_units,
)
# third pass: convert other arguments
for ndx in unit_args_ndx:
for pos in indices[ndx]:
if isinstance(values[pos], ureg.Quantity):
values[pos] = ureg._convert(
values[pos]._magnitude, values[pos]._units, args_as_uc[ndx][0]
)
else:
if strict:
if isinstance(values[pos], str):
# if the value is a string, we try to parse it
tmp_value = ureg.parse_expression(values[pos])
values[pos] = ureg._convert(
tmp_value._magnitude,
tmp_value._units,
args_as_uc[ndx][0],
)
else:
raise ValueError(
"A wrapped function using strict=True requires "
"quantity or a string for all arguments with not None units. "
"(error found for {}, {})".format(
args_as_uc[ndx][0], values[pos]
)
)
if has_varargs:
ndx = 0
for name, param in sig.parameters.items():
if param.kind == Parameter.VAR_KEYWORD:
continue
positions = indices[ndx]
if param.kind == Parameter.VAR_POSITIONAL:
bound.arguments[name] = tuple(values[pos] for pos in positions)
else:
bound.arguments[name] = values[positions[0]]
ndx += 1
return bound.args, bound.kwargs, values_by_name
# unpack kwargs
for i, (param_name, param) in enumerate(sig.parameters.items()):
if param.kind == Parameter.VAR_KEYWORD:
continue
if i >= len_initial_values:
kw[param_name] = values[i]
return values[:len_initial_values], kw, values_by_name
return _converter
def _apply_defaults(sig, args, kwargs):
"""Apply default keyword arguments.
Named keywords may have been left blank. This function applies the default
values so that every argument is defined.
"""
for i, param in enumerate(sig.parameters.values()):
if (
i >= len(args)
and param.default != Parameter.empty
and param.name not in kwargs
):
kwargs[param.name] = param.default
return list(args), kwargs
@overload
def wraps[QuantityT: Quantity, UnitT: Unit](
ureg: GenericPlainRegistry[QuantityT, UnitT],
ret: str | Unit,
args: str | Unit | Iterable[str | Unit | None] | None,
strict: bool = True,
) -> Callable[[Callable[..., Any]], Callable[..., QuantityT]]: ...
# Skipped conversion and multiple return values need not produce a Quantity.
@overload
def wraps[QuantityT: Quantity, UnitT: Unit](
ureg: GenericPlainRegistry[QuantityT, UnitT],
ret: str | Unit | Iterable[str | Unit | None] | None,
args: str | Unit | Iterable[str | Unit | None] | None,
strict: bool = True,
) -> Callable[[Callable[..., Any]], Callable[..., object]]: ...
[docs]
def wraps[QuantityT: Quantity, UnitT: Unit](
ureg: GenericPlainRegistry[QuantityT, UnitT],
ret: str | Unit | Iterable[str | Unit | None] | None,
args: str | Unit | Iterable[str | Unit | None] | None,
strict: bool = True,
) -> Callable[[Callable[..., Any]], Callable[..., object]]:
"""Wraps a function to become pint-aware.
Use it when a function requires a numerical value but in some specific
units. The wrapper function will take a pint quantity, convert to the units
specified in `args` and then call the wrapped function with the resulting
magnitude.
The value returned by the wrapped function will be converted to the units
specified in `ret`.
Parameters
----------
ureg : pint.UnitRegistry
a UnitRegistry instance.
ret : str, pint.Unit, or iterable of str or pint.Unit
Units of each of the return values. Use `None` to skip argument conversion.
args : str, pint.Unit, or iterable of str or pint.Unit
Units of each of the input arguments. Use `None` to skip argument conversion.
One unit specification applies to every element of a ``*args`` parameter.
If that specification defines a reference such as ``"=A"``, the first
element defines its unit; an empty ``*args`` defines it as dimensionless.
strict : bool
Indicates that only quantities are accepted. (Default value = True)
Returns
-------
callable
the wrapper function.
Raises
------
TypeError
if the number of given arguments does not match the number of function parameters.
if any of the provided arguments is not a unit a string or Quantity
"""
if not isinstance(args, (list, tuple)):
args = (args,)
for arg in args:
if arg is not None and not isinstance(arg, (ureg.Unit, str)):
raise TypeError(
"wraps arguments must by of type str or Unit, not %s (%s)"
% (type(arg), arg)
)
is_ret_container = isinstance(ret, (list, tuple))
if is_ret_container:
for arg in ret:
if arg is not None and not isinstance(arg, (ureg.Unit, str)):
raise TypeError(
"wraps 'ret' argument must by of type str or Unit, not %s (%s)"
% (type(arg), arg)
)
ret = ret.__class__([_to_units_container(arg, ureg) for arg in ret])
else:
if ret is not None and not isinstance(ret, (ureg.Unit, str)):
raise TypeError(
"wraps 'ret' argument must by of type str or Unit, not %s (%s)"
% (type(ret), ret)
)
ret = _to_units_container(ret, ureg)
def decorator(func: Callable[..., Any]) -> Callable[..., object]:
sig = signature(func)
params = tuple(sig.parameters.values())
has_varargs = any(param.kind == Parameter.VAR_POSITIONAL for param in params)
has_var_keyword = any(param.kind == Parameter.VAR_KEYWORD for param in params)
count_params = sum(param.kind != Parameter.VAR_KEYWORD for param in params)
converter_args = args
if has_var_keyword and len(args) == count_params + 1 and args[-1] is None:
converter_args = args[:-1]
if len(converter_args) != count_params:
raise TypeError(
"%s takes %i parameters, but %i units were passed"
% (func.__name__, count_params, len(args))
)
converter = _parse_wrap_args(converter_args, ureg, has_varargs=has_varargs)
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, **kw) -> object:
if not has_varargs:
values, kw = _apply_defaults(sig, values, kw)
# In principle, the values are used as is
# When then extract the magnitudes when needed.
new_values, new_kw, values_by_name = converter(
ureg, sig, values, kw, strict
)
result = func(*new_values, **new_kw)
if is_ret_container:
out_units = (
_replace_units(r, values_by_name) if is_ref else r
for (r, is_ref) in ret
)
return ret.__class__(
res if unit is None else ureg.Quantity(res, unit)
for unit, res in zip_longest(out_units, result)
)
if ret[0] is None:
return result
return ureg.Quantity(
result, _replace_units(ret[0], values_by_name) if ret[1] else ret[0]
)
return wrapper
return decorator
[docs]
def check[F: FuncType](
ureg: UnitRegistry, *args: str | UnitsContainer | Unit | None
) -> Callable[[F], F]:
"""Decorator to for quantity type checking for function inputs.
Use it to ensure that the decorated function input parameters match
the expected dimension of pint quantity.
The wrapper function raises:
- `pint.DimensionalityError` if an argument doesn't match the required dimensions.
ureg : UnitRegistry
a UnitRegistry instance.
args : str or UnitContainer or None
Dimensions of each of the input arguments, one per parameter of the
decorated function in signature order. A dimension given for a
``*args`` parameter is checked against each value supplied through it.
Use `None` to skip checking a parameter; ``**kwargs`` requires `None`.
Returns
-------
callable
the wrapped function.
Raises
------
TypeError
If the number of given dimensions does not match the number of function
parameters.
ValueError
If the any of the provided dimensions cannot be parsed as a dimension.
"""
dimensions = [
ureg.get_dimensionality(dim) if dim is not None else None for dim in args
]
def decorator(func):
sig = signature(func)
params = tuple(sig.parameters.values())
count_params = len(params)
if len(dimensions) != count_params:
raise TypeError(
"%s takes %i parameters, but %i dimensions were passed"
% (func.__name__, count_params, len(dimensions))
)
for param, dim in zip(params, dimensions):
if param.kind is Parameter.VAR_KEYWORD and dim is not None:
raise TypeError("use None for the **kwargs dimension")
needs_binding = any(
param.kind
not in (Parameter.POSITIONAL_ONLY, Parameter.POSITIONAL_OR_KEYWORD)
for param in params
)
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(*args, **kwargs):
if needs_binding:
# Binding distinguishes positional values from keyword-only and
# variadic parameters before dimensional checks run.
bound = sig.bind(*args, **kwargs)
bound.apply_defaults()
values = [bound.arguments[param.name] for param in params]
else:
values = list(args)
used_kwargs = 0
for param in params[len(args) :]:
if (
param.name in kwargs
and param.kind is not Parameter.POSITIONAL_ONLY
):
values.append(kwargs[param.name])
used_kwargs += 1
elif param.default is not Parameter.empty:
values.append(param.default)
else:
break
if len(values) != count_params or used_kwargs != len(kwargs):
# Invalid call: let it raise TypeError before any checks.
return func(*args, **kwargs)
for param, dim, value in zip(params, dimensions, values):
if dim is None:
continue
if param.kind is Parameter.VAR_POSITIONAL:
items = value
else:
items = (value,)
for item in items:
if not ureg.Quantity(item).check(dim):
val_dim = ureg.get_dimensionality(item)
raise DimensionalityError(item, "a quantity of", val_dim, dim)
return func(*args, **kwargs)
return wrapper
return decorator
