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
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
|
"""
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
|
