From ad42b41919b068c6b9cb488a8b719f0cf8d1cac7 Mon Sep 17 00:00:00 2001 From: Hernan Grecco Date: Sat, 20 Jan 2024 02:52:36 -0300 Subject: [PATCH] Added class to enable dynamic registration (backwards compatiblity) --- pint/delegates/formatter/_to_register.py | 117 +++++++++++++++++++++++ 1 file changed, 117 insertions(+) create mode 100644 pint/delegates/formatter/_to_register.py diff --git a/pint/delegates/formatter/_to_register.py b/pint/delegates/formatter/_to_register.py new file mode 100644 index 000000000..778733801 --- /dev/null +++ b/pint/delegates/formatter/_to_register.py @@ -0,0 +1,117 @@ +""" + pint.delegates.formatter.base_formatter + ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + Common class and function for all formatters. + :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, Callable +from ...compat import ndarray, np, Unpack +from ._helpers import ( + split_format, + join_mu, +) + +from ..._typing import Magnitude + +from ._unit_handlers import format_compound_unit, BabelKwds, override_locale + +if TYPE_CHECKING: + from ...facets.plain import PlainQuantity, PlainUnit, MagnitudeT + from ...registry import UnitRegistry + + +def register_unit_format(name: str): + """register a function as a new format for units + + The registered function must have a signature of: + + .. code:: python + + def new_format(unit, registry, **options): + pass + + Parameters + ---------- + name : str + The name of the new format (to be used in the format mini-language). A error is + raised if the new format would overwrite a existing format. + + Examples + -------- + .. code:: python + + @pint.register_unit_format("custom") + def format_custom(unit, registry, **options): + result = "" # do the formatting + return result + + + ureg = pint.UnitRegistry() + u = ureg.m / ureg.s ** 2 + f"{u:custom}" + """ + + from ...formatting import _ORPHAN_FORMATTER + + # TODO: kwargs missing in typing + def wrapper(func: Callable[[PlainUnit, UnitRegistry], str]): + if name in _ORPHAN_FORMATTER._formatters: + raise ValueError(f"format {name!r} already exists") # or warn instead + + class NewFormatter: + def format_magnitude( + self, + magnitude: Magnitude, + mspec: str = "", + **babel_kwds: Unpack[BabelKwds], + ) -> str: + with override_locale( + mspec, babel_kwds.get("locale", None) + ) as format_number: + if isinstance(magnitude, ndarray) and magnitude.ndim > 0: + # Use custom ndarray text formatting--need to handle scalars differently + # since they don't respond to printoptions + with np.printoptions(formatter={"float_kind": format_number}): + mstr = format(magnitude).replace("\n", "") + else: + mstr = format_number(magnitude) + + return mstr + + def format_unit( + self, unit: PlainUnit, uspec: str = "", **babel_kwds: Unpack[BabelKwds] + ) -> str: + units = unit._REGISTRY.UnitsContainer( + format_compound_unit(unit, uspec, **babel_kwds) + ) + + return func(units, registry=unit._REGISTRY, **babel_kwds) + + def format_quantity( + self, + quantity: PlainQuantity[MagnitudeT], + qspec: str = "", + **babel_kwds: Unpack[BabelKwds], + ) -> str: + registry = quantity._REGISTRY + + mspec, uspec = split_format( + qspec, + registry.formatter.default_format, + registry.separate_format_defaults, + ) + + joint_fstring = "{} {}" + return join_mu( + joint_fstring, + self.format_magnitude(quantity.magnitude, mspec, **babel_kwds), + self.format_unit(quantity.units, uspec, **babel_kwds), + ) + + _ORPHAN_FORMATTER._formatters[name] = NewFormatter() + + return wrapper