eng-lab/hwe/units.py

"""Units for the whole hardware engine: one registry, so no number is ever unitless.

Why this file exists (EL01, section 2.2)
    A bare number is a bug waiting to happen: is 1300 a mass in grams, a capacity
    in mAh, or a thrust in gram-force? Every quantity in the engine is therefore a
    pint Quantity: a magnitude and a unit travelling together. pint refuses to add
    metres to seconds, converts grams-force to newtons on request, and raises
    DimensionalityError the moment an equation stops making physical sense.

    pint cannot mix quantities made by two different registries, so this module
    owns the ONLY one. Import `ureg`, `Q_`, and the helpers from here; never call
    pint.UnitRegistry() anywhere else.

Everyday use
    Q_(9.81, "m/s^2")          a quantity from a number and a unit
    Q_("1300 mA*h")            the same from one string (design files use this)
    thrust.to("N")             convert; raises DimensionalityError if impossible
    length.m_as("mm")          the bare number in a stated unit. This is the ONE
                               sanctioned way to leave the unit system, and only
                               at a named boundary (for example, build123d works
                               in plain millimetres).
    f"{q:.3~P}"                pretty print with 3 significant digits: "9.81 m/s²"

The @quantities decorator
    Every function you write in hwe/core/ declares the dimension of each input
    and of its result. The decorator checks them on every call, so a bare float
    or a wrong unit fails loudly at the function boundary instead of silently
    corrupting a result three analyses later.
"""
from __future__ import annotations

import functools
import inspect

import pint

# The one registry. Offset units (degC) auto-convert to kelvin in arithmetic so
# that temperature maths cannot silently go wrong in later modules.
ureg = pint.UnitRegistry(autoconvert_offset_to_baseunit=True)
ureg.formatter.default_format = "~P"      # abbreviated, pretty: "9.81 m/s²"
ureg.formatter.default_sort_func = None   # keep units in written order: "W⋅h", not "h⋅W"
pint.set_application_registry(ureg)

Q_ = ureg.Quantity
Quantity = ureg.Quantity                    # for type hints: x: Quantity
DimensionalityError = pint.DimensionalityError

# Standard gravity g0. Exact by definition (3rd CGPM, 1901). Local gravity differs
# by a few tenths of a percent (about 9.780 m/s² at the equator, 9.832 m/s² at the
# poles); the design file can override it for a specific site.
G0 = Q_(9.80665, "m/s^2")


def is_quantity(x) -> bool:
    """True if x is a Quantity from THIS registry."""
    return isinstance(x, Quantity)


def require(x, dimension: str, name: str = "value") -> Quantity:
    """Return x unchanged if it is a Quantity with the given dimension; raise otherwise.

    dimension is a pint dimensionality string such as "[mass]", "[length]",
    "[mass] * [length] / [time] ** 2" (a force), or "[]" for a pure number.

    Raises TypeError for a bare number (the most common beginner slip) and
    DimensionalityError for a quantity of the wrong kind (grams where newtons
    belong, for example).
    """
    if not is_quantity(x):
        raise TypeError(
            f"{name} must carry units, e.g. Q_(1.2, 'kg'); got the bare {type(x).__name__} {x!r}"
        )
    wanted = ureg.get_dimensionality(dimension)
    if x.dimensionality != wanted:
        raise DimensionalityError(x.units, dimension, extra_msg=f" ({name})")
    return x


def quantities(returns: str | None = None, **dims: str):
    """Decorator: check the dimension of named arguments and of the return value.

    Example (this is how every starter in hwe/core/ is declared):

        @quantities(mass="[mass]", returns="[force]")
        def weight(mass):
            ...

    Arguments not named in **dims are not checked (counts, fractions, names).
    Dimensionless results ("[]") are converted to plain dimensionless form, so a
    ratio built from gram-force and newtons reads as the pure number it is.
    """
    def decorate(func):
        signature = inspect.signature(func)

        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            bound = signature.bind(*args, **kwargs)
            bound.apply_defaults()
            for arg_name, dim in dims.items():
                value = bound.arguments[arg_name]
                if isinstance(value, (list, tuple)):
                    for i, item in enumerate(value):
                        require(item, dim, f"{func.__name__}(): {arg_name}[{i}]")
                else:
                    require(value, dim, f"{func.__name__}(): {arg_name}")
            result = func(*args, **kwargs)
            if returns is not None:
                require(result, returns, f"{func.__name__}() result")
                if returns == "[]":
                    result = result.to("dimensionless")
            return result

        wrapper.dims = dict(dims)
        wrapper.returns = returns
        return wrapper
    return decorate


def parse(text_or_number, default_unit: str | None = None) -> Quantity:
    """Turn a design-file entry such as "150 mm" into a Quantity.

    A bare number is accepted only when default_unit is given (counts and pure
    ratios); anything else must spell out its unit.
    """
    if is_quantity(text_or_number):
        return text_or_number
    if isinstance(text_or_number, (int, float)):
        if default_unit is None:
            raise TypeError(f"{text_or_number!r} has no unit; write it as a string such as '150 mm'")
        return Q_(text_or_number, default_unit)
    return Q_(str(text_or_number))


def _number(x: float, digits: int) -> str:
    """x with `digits` significant digits; plain notation from 0.001 to a million."""
    import math
    if x == 0 or not math.isfinite(x):
        return f"{x:g}"
    if 1e-3 <= abs(x) < 1e6:
        decimals = max(0, digits - 1 - math.floor(math.log10(abs(x))))
        return f"{x:,.{decimals}f}"
    return f"{x:.{digits}g}"


def fmt(q: Quantity, digits: int = 3) -> str:
    """A quantity as short text with the given significant digits: '6.37 N', '18,547 mm³'."""
    if not is_quantity(q):
        return _number(float(q), digits)
    unit = f"{q.units:~P}"
    number = _number(float(q.m), digits)
    return f"{number} {unit}" if unit else number