diff --git a/pint/facets/plain/quantity.py b/pint/facets/plain/quantity.py index f7710df97..199b19725 100644 --- a/pint/facets/plain/quantity.py +++ b/pint/facets/plain/quantity.py @@ -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: diff --git a/pint/facets/plain/registry.py b/pint/facets/plain/registry.py index 277a6f7a2..b09c0f1f9 100644 --- a/pint/facets/plain/registry.py +++ b/pint/facets/plain/registry.py @@ -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]] = {} @@ -135,6 +138,7 @@ def __eq__(self, other: Any): return False attrs = ( "dimensional_equivalents", + "kind_dimensional_equivalents", "root_units", "dimensionality", "parse_unit", @@ -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: @@ -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 diff --git a/pint/testsuite/test_kind.py b/pint/testsuite/test_kind.py index 9e6ecd801..3d4604b06 100644 --- a/pint/testsuite/test_kind.py +++ b/pint/testsuite/test_kind.py @@ -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()