Source code for genno.core.quantity

import sys
from typing import TYPE_CHECKING, Literal, Type, Union

from .attrseries import AttrSeries
from .base import BaseQuantity
from .sparsedataarray import SparseDataArray

if TYPE_CHECKING:
    from typing import TypeAlias


[docs] def assert_quantity(*args): """Assert that each of `args` is a Quantity object. Raises ------ TypeError with a indicative message. """ for i, arg in enumerate(args): if not isinstance(arg, BaseQuantity): raise TypeError( f"arg #{i+1} ({repr(arg)[:20]}) is not Quantity; likely an incorrect key" )
[docs] def get_class() -> Type[Union[AttrSeries, SparseDataArray]]: """Get the current :class:`.Quantity` implementation in use. Returns one of the classes :class:`.AttrSeries` or :class:`.SparseDataArray`. """ global Quantity return Quantity
[docs] def set_class( name: Literal["AttrSeries", "SparseDataArray"] = "AttrSeries", ) -> Type[Union[AttrSeries, SparseDataArray]]: """Set the :class:`.Quantity` implementation to be used. This also updates :py:`genno.Quantity` and :py:`genno.quantity.Quantity` to refer to the selected class. It does **not** update previously-imported references to one class or the other; code that uses :func:`.set_class` should refer to one of those two locations: .. code-block:: import genno from genno import Quantity # AttrSeries, by default Quantity() # AttrSeries genno.Quantity() # AttrSeries genno.set_class("SparseDataArray") Quantity() # AttrSeries genno.Quantity() # SparseDataArray Another approach is to update the local reference with the return value of the function: .. code-block:: python from genno import Quantity, set_class Quantity() # AttrSeries Quantity = set_class("SparseDataArray") Quantity() # SparseDataArray In code that does not use :func:`.set_class`, :py:`from genno import Quantity` is safe. See also -------- .AnyQuantity """ global CLASS try: cls = {"AttrSeries": AttrSeries, "SparseDataArray": SparseDataArray}[name] except KeyError: raise ValueError(f"no Quantity implementation {name}") # Update globals in the current module and at the top level for module in "genno.core.quantity", "genno": setattr(sys.modules[module], "Quantity", cls) # Update global CLASS = cls.__name__ return cls
#: Class used to implement :class:`.Quantity`. Quantity: "TypeAlias" = AttrSeries #: Name of :class:`.Quantity`. CLASS = "AttrSeries" #: Either :class:`.AttrSeries` or :class:`.SparseDataArray`. Code in :mod:`genno` or #: user code that receives or returns any Quantity implementation should be typed with #: this type. AnyQuantity = Union[AttrSeries, SparseDataArray]