Source code for pinttr._generator

from contextlib import contextmanager
from copy import copy
from typing import Callable, Union

import attrs
import pint


[docs] @attrs.define class UnitGenerator: """ A callable object which returns units objects. Stored units can be contextually overridden using the :meth:`~pinttrs.UnitGenerator.override` method. .. seealso:: :class:`~pinttrs.UnitContext` :Attributes / constructor arguments: **units** (:class:`pint.Unit` or :class:`Callable`) – Stored units or generator. """ units: Union[pint.Unit, Callable] = attrs.field()
[docs] def __call__(self) -> pint.Unit: """ :returns: If ``units`` is a :class:`pint.Unit`, it is returned; if ``units`` is a callable (typically, another :class:`~pinttrs.UnitGenerator`), the result of its evaluation will be returned. """ if callable(self.units): return self.units() return self.units
[docs] @contextmanager def override(self, units: Union[pint.Unit, Callable, str]) -> None: """ Temporarily override the value of ``units``. The initial value of ``units`` is restored upon leaving context. :param units: Temporary replacement for ``units``. String values are interpreted based on the unit registry of currently stored units. """ units_old = copy(self.units) if isinstance(units, str): # Safeguard to convert strings if callable(self.units): self.units = self.units()._REGISTRY.Unit(units) else: self.units = self.units._REGISTRY.Unit(units) else: self.units = units try: yield finally: self.units = units_old