diff --git a/src/ansys/units/__init__.py b/src/ansys/units/__init__.py index df77c80e..221c14a5 100644 --- a/src/ansys/units/__init__.py +++ b/src/ansys/units/__init__.py @@ -43,7 +43,7 @@ from ansys.units.quantity import Quantity, get_si_value # noqa: F401 from ansys.units.systems import UnitSystem # noqa: F401 from ansys.units.unit import Unit # noqa: F401 -from ansys.units.unit_registry import UnitRegistry # noqa: F401 +from ansys.units.unit_registry import UnitRegistry, register_unit # noqa: F401 _THIS_DIRNAME = os.path.dirname(__file__) _README_FILE = os.path.normpath(os.path.join(_THIS_DIRNAME, "docs", "README.rst")) diff --git a/src/ansys/units/unit_registry.py b/src/ansys/units/unit_registry.py index 80a86607..01b2dcbd 100644 --- a/src/ansys/units/unit_registry.py +++ b/src/ansys/units/unit_registry.py @@ -25,7 +25,41 @@ import yaml -from ansys.units import Unit +from ansys.units import Unit, _base_units, _derived_units + + +def register_unit(unit: str, composition: str, factor): + """ + Register a customized unit. + + Parameters + ---------- + unit: str + Name of the unit + composition: str + String chain of combined units + factor: float + Scaling factor + + Examples + -------- + >>> from ansys.units import register_unit + >>> register_unit(unit="Q", composition="N m", factor=1) + >>> ureg = UnitRegistry() + >>> ureg.Q + _name: Q + _dimensions: {'MASS': 1.0, 'LENGTH': 2.0, 'TIME': -2.0} + _composition: N m + _factor: 1 + _si_units: kg m^2 s^-2 + _si_scaling_factor: 1.0 + _si_offset: 0.0 + """ + if unit in _base_units or unit in _derived_units: + raise UnitAlreadyRegistered(unit) + else: + register = {unit: {"composition": composition, "factor": factor}} + _derived_units.update(register) class UnitRegistry: @@ -55,6 +89,8 @@ class UnitRegistry: def __init__(self, config="cfg.yaml", other: dict = None): unitdict = other or {} + global _derived_units + unitdict.update(_derived_units) if config: file_dir = os.path.dirname(__file__) @@ -62,10 +98,10 @@ def __init__(self, config="cfg.yaml", other: dict = None): with open(qc_path, "r") as qc_yaml: qc_data = yaml.safe_load(qc_yaml) - _base_units: dict = qc_data["base_units"] - _derived_units: dict = qc_data["derived_units"] + _qc_base_units: dict = qc_data["base_units"] + _qc_derived_units: dict = qc_data["derived_units"] - unitdict.update(**_base_units, **_derived_units) + unitdict.update(**_qc_base_units, **_qc_derived_units) for unit in unitdict: setattr(self, unit, Unit(unit, unitdict[unit])) diff --git a/tests/test_unit_registry.py b/tests/test_unit_registry.py index 1b684305..70fbc8a3 100644 --- a/tests/test_unit_registry.py +++ b/tests/test_unit_registry.py @@ -25,7 +25,7 @@ import pytest -from ansys.units import Unit, UnitRegistry, _base_units +from ansys.units import Unit, UnitRegistry, _base_units, register_unit from ansys.units.unit_registry import UnitAlreadyRegistered @@ -97,6 +97,22 @@ def test_immutability(): ur.m = Unit("ft") +def test_register_unit(): + register_unit(unit="Q", composition="N m", factor=1) + ur = UnitRegistry() + assert ur.Q == ur.J + + with pytest.raises(UnitAlreadyRegistered): + register_unit(unit="Q", composition="N m", factor=1) + + register_unit(unit="Z", composition="N m", factor=1) + with pytest.raises(AttributeError): + ur.Z + + ur = UnitRegistry() + assert ur.Q == ur.J == ur.Z + + def test_error_message(): e1 = UnitAlreadyRegistered("kg") expected_str = "Unable to override `kg` it has already been registered."