Skip to content

Commit

Permalink
compatable
Browse files Browse the repository at this point in the history
  • Loading branch information
andrewgsavage committed Apr 11, 2024
1 parent ba366c2 commit ef44882
Show file tree
Hide file tree
Showing 3 changed files with 26 additions and 0 deletions.
3 changes: 3 additions & 0 deletions pint/facets/plain/quantity.py
Original file line number Diff line number Diff line change
Expand Up @@ -461,6 +461,9 @@ def to_kind(self, kind):
result.kind = kind
return result

def compatable_kinds(self):
return self._REGISTRY.get_compatible_kinds(self.dimensionality)

def is_compatible_with(
self, other: Any, *contexts: str | Context, **ctx_kwargs: Any
) -> bool:
Expand Down
16 changes: 16 additions & 0 deletions pint/facets/plain/registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -116,6 +116,9 @@ def __init__(self) -> None:
#: Maps dimensionality (UnitsContainer) to Units (str)
self.dimensional_equivalents: dict[UnitsContainer, frozenset[str]] = {}

#: Maps dimensionality (UnitsContainer) to kinds (str)
self.kind_dimensional_equivalents: dict[UnitsContainer, frozenset[str]] = {}

#: Maps dimensionality (UnitsContainer) to Dimensionality (UnitsContainer)
# TODO: this description is not right.
self.root_units: dict[UnitsContainer, tuple[Scalar, UnitsContainer]] = {}
Expand All @@ -135,6 +138,7 @@ def __eq__(self, other: Any):
return False
attrs = (
"dimensional_equivalents",
"kind_dimensional_equivalents",
"root_units",
"dimensionality",
"parse_unit",
Expand Down Expand Up @@ -640,6 +644,13 @@ def _build_cache(self, loaded_files=None) -> None:

except Exception as exc:
logger.warning(f"Could not resolve {unit_name}: {exc!r}")

for kind_name in self._dimensions:
di = self._get_dimensionality(self.UnitsContainer({kind_name: 1}))
dimeq_set = self._cache.kind_dimensional_equivalents.setdefault(
di, set()
)
dimeq_set.add(kind_name)
return self._cache

def get_name(self, name_or_alias: str, case_sensitive: bool | None = None) -> str:
Expand Down Expand Up @@ -967,6 +978,11 @@ def _get_compatible_units(
src_dim = self._get_dimensionality(input_units)
return self._cache.dimensional_equivalents.setdefault(src_dim, frozenset())

def get_compatible_kinds(self, dimensionality: UnitsContainer) -> frozenset[str]:
return self._cache.kind_dimensional_equivalents.setdefault(
dimensionality, frozenset()
)

# TODO: remove context from here
def is_compatible_with(
self, obj1: Any, obj2: Any, *contexts: str | Context, **ctx_kwargs
Expand Down
7 changes: 7 additions & 0 deletions pint/testsuite/test_kind.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,3 +32,10 @@ def test_torque_energy(self):
# Torque is not energy so cannot be added
with pytest.raises(ValueError):
energy + torque

def test_compatable_kinds(self):
ureg = UnitRegistry()
Q_ = ureg.Quantity
q = Q_(1, "N m")
assert "[torque]" in q.compatable_kinds()
assert "[energy]" in q.compatable_kinds()

0 comments on commit ef44882

Please sign in to comment.