Skip to main content
`Ctrl`+`K`

#

# Source code for pint.facets.group.registry

    """
    pint.facets.group.registry
    ~~~~~~~~~~~~~~~~~~~~~~~~~~

    :copyright: 2022 by Pint Authors, see AUTHORS for more details.
    :license: BSD, see LICENSE for more details.
    """

    from __future__ import annotations

    from typing import TYPE_CHECKING, Any, Generic

    from ... import errors
    from ...compat import TypeAlias

    if TYPE_CHECKING:
        from ..._typing import Unit, UnitsContainer

    from ...util import create_class_with_registry, to_units_container
    from ..plain import (
        GenericPlainRegistry,
        QuantityT,
        UnitDefinition,
        UnitT,
    )
    from . import objects
    from .definitions import GroupDefinition

    class GenericGroupRegistry(
        Generic[QuantityT, UnitT], GenericPlainRegistry[QuantityT, UnitT]
    ):
        """Handle of Groups.

        Group units

        Capabilities:
        - Register groups.
        - Parse @group directive.
        """

        # TODO: Change this to Group: Group to specify class
        # and use introspection to get system class as a way
        # to enjoy typing goodies
        Group = type[objects.Group]

        def __init__(self, **kwargs):
            super().__init__(**kwargs)
            #: Map group name to group.
            self._groups: dict[str, objects.Group] = {}
            self._groups["root"] = self.Group("root")

        def _init_dynamic_classes(self) -> None:
            """Generate subclasses on the fly and attach them to self"""
            super()._init_dynamic_classes()
            self.Group = create_class_with_registry(self, objects.Group)

        def _after_init(self) -> None:
            """Invoked at the end of ``__init__``.

            - Create default group and add all orphan units to it
            - Set default system
            """
            super()._after_init()

            #: Copy units not defined in any group to the default group
            if "group" in self._defaults:
                grp = self.get_group(self._defaults["group"], True)
                group_units = frozenset(
                    [
                        member
                        for group in self._groups.values()
                        if group.name != "root"
                        for member in group.members
                    ]
                )
                all_units = self.get_group("root", False).members
                grp.add_units(*(all_units - group_units))

        def _register_definition_adders(self) -> None:
            super()._register_definition_adders()
            self._register_adder(GroupDefinition, self._add_group)

        def _add_unit(self, definition: UnitDefinition):
            super()._add_unit(definition)
            # TODO: delta units are missing
            self.get_group("root").add_units(definition.name)

        def _add_group(self, gd: GroupDefinition):
            if gd.name in self._groups:
                raise ValueError(f"Group {gd.name} already present in registry")
            try:
                # As a Group is a SharedRegistryObject
                # it adds itself to the registry.
                self.Group.from_definition(gd)
            except KeyError as e:
                raise errors.DefinitionSyntaxError(f"unknown dimension {e} in context")

        def get_group(self, name: str, create_if_needed: bool = True) -> objects.Group:
            """Return a Group.

            Parameters
            ----------
            name : str
                Name of the group to be
            create_if_needed : bool
                If True, create a group if not found. If False, raise an Exception.
                (Default value = True)

            Returns
            -------
            Group
                Group
            """
            if name in self._groups:
                return self._groups[name]

            if not create_if_needed:
                raise ValueError("Unknown group %s" % name)

            return self.Group(name)

        def get_compatible_units(
            self, input_units: UnitsContainer, group: str | None = None
        ) -> frozenset[Unit]:
            """ """
            if group is None:
                return super().get_compatible_units(input_units)

            input_units = to_units_container(input_units)

            equiv = self._get_compatible_units(input_units, group)

            return frozenset(self.Unit(eq) for eq in equiv)

        def _get_compatible_units(
            self, input_units: UnitsContainer, group: str | None = None
        ) -> frozenset[str]:
            ret = super()._get_compatible_units(input_units)

            if not group:
                return ret

            if group in self._groups:
                members = self._groups[group].members
            else:
                raise ValueError("Unknown Group with name '%s'" % group)
            return frozenset(ret & members)

    class GroupRegistry(
        GenericGroupRegistry[objects.GroupQuantity[Any], objects.GroupUnit]
    ):
        Quantity: TypeAlias = objects.GroupQuantity[Any]
        Unit: TypeAlias = objects.GroupUnit