summaryrefslogtreecommitdiff
path: root/pint/facets/group/registry.py
blob: c4ed0be2ee6cb11cb32c054bfb15220c4ca256f8 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
"""
    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, Dict, FrozenSet

from ... import errors

if TYPE_CHECKING:
    from pint import Unit

from ...util import build_dependent_class, create_class_with_registry
from ..plain import PlainRegistry, UnitDefinition
from .definitions import GroupDefinition
from .objects import Group


class GroupRegistry(PlainRegistry):
    """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_class = Group

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

    def __init_subclass__(cls, **kwargs):
        super().__init_subclass__()
        cls.Group = build_dependent_class(cls, "Group", "_group_class")

    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, self.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) -> 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, group) -> FrozenSet["Unit"]:

        ret = super()._get_compatible_units(input_units, group)

        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)