custom unit formatter (#278)

Author: keewis

cf_xarray/tests/test_units.py

@@ -8,53 +8,68 @@
 
 pytest.importorskip("pint")
 
-from ..units import units
+from ..units import units as ureg
 
 
 def test_added_degrees_units():
     """Test that our added degrees units are present in the registry."""
     # Test equivalence of abbreviations/aliases to our defined names
-    assert str(units("degrees_N").units) == "degrees_north"
-    assert str(units("degreesN").units) == "degrees_north"
-    assert str(units("degree_north").units) == "degrees_north"
-    assert str(units("degree_N").units) == "degrees_north"
-    assert str(units("degreeN").units) == "degrees_north"
-    assert str(units("degrees_E").units) == "degrees_east"
-    assert str(units("degreesE").units) == "degrees_east"
-    assert str(units("degree_east").units) == "degrees_east"
-    assert str(units("degree_E").units) == "degrees_east"
-    assert str(units("degreeE").units) == "degrees_east"
+    assert str(ureg("degrees_N").units) == "degrees_north"
+    assert str(ureg("degreesN").units) == "degrees_north"
+    assert str(ureg("degree_north").units) == "degrees_north"
+    assert str(ureg("degree_N").units) == "degrees_north"
+    assert str(ureg("degreeN").units) == "degrees_north"
+    assert str(ureg("degrees_E").units) == "degrees_east"
+    assert str(ureg("degreesE").units) == "degrees_east"
+    assert str(ureg("degree_east").units) == "degrees_east"
+    assert str(ureg("degree_E").units) == "degrees_east"
+    assert str(ureg("degreeE").units) == "degrees_east"
 
     # Test equivalence of our defined units to base units
-    assert units("degrees_north") == units("degrees")
-    assert units("degrees_north").to_base_units().units == units.radian
-    assert units("degrees_east") == units("degrees")
-    assert units("degrees_east").to_base_units().units == units.radian
+    assert ureg("degrees_north") == ureg("degrees")
+    assert ureg("degrees_north").to_base_units().units == ureg.radian
+    assert ureg("degrees_east") == ureg("degrees")
+    assert ureg("degrees_east").to_base_units().units == ureg.radian
 
 
 def test_gpm_unit():
     """Test that the gpm unit does alias to meters."""
-    x = 1 * units("gpm")
+    x = 1 * ureg("gpm")
     assert str(x.units) == "meter"
 
 
 def test_psu_unit():
     """Test that the psu unit are present in the registry."""
-    x = 1 * units("psu")
+    x = 1 * ureg("psu")
     assert str(x.units) == "practical_salinity_unit"
 
 
 def test_percent_units():
     """Test that percent sign units are properly parsed and interpreted."""
-    assert str(units("%").units) == "percent"
+    assert str(ureg("%").units) == "percent"
 
 
 @pytest.mark.xfail(reason="not supported by pint, yet: hgrecco/pint#1295")
 def test_udunits_power_syntax():
     """Test that UDUNITS style powers are properly parsed and interpreted."""
-    assert units("m2 s-2").units == units.m ** 2 / units.s ** 2
+    assert ureg("m2 s-2").units == ureg.m ** 2 / ureg.s ** 2
 
 
 def test_udunits_power_syntax_parse_units():
     """Test that UDUNITS style powers are properly parsed and interpreted."""
-    assert units.parse_units("m2 s-2") == units.m ** 2 / units.s ** 2
+    assert ureg.parse_units("m2 s-2") == ureg.m ** 2 / ureg.s ** 2
+
+
+@pytest.mark.parametrize(
+    ["units", "expected"],
+    (
+        ("kg ** 2", "kg2"),
+        ("m ** -1", "m-1"),
+        ("m ** 2 / s ** 2", "m2 s-2"),
+        ("m ** 3 / (kg * s ** 2)", "m3 kg-1 s-2"),
+    ),
+)
+def test_udunits_format(units, expected):
+    u = ureg.parse_units(units)
+
+    assert f"{u:cf}" == expected

cf_xarray/units.py

@@ -1,8 +1,4 @@
-r"""Module to provide unit support via pint approximating UDUNITS/CF.
-
-Reused with modification from MetPy under the terms of the BSD 3-Clause License.
-Copyright (c) 2015,2017,2019 MetPy Developers.
-"""
+"""Module to provide unit support via pint approximating UDUNITS/CF."""
 import functools
 import re
 import warnings
@@ -14,6 +10,59 @@
     UnitStrippedWarning,
 )
 
+# from `xclim`'s unit support module with permission of the maintainers
+try:
+
+    @pint.register_unit_format("cf")
+    def short_formatter(unit, registry, **options):
+        """Return a CF-compliant unit string from a `pint` unit.
+
+        Parameters
+        ----------
+        unit : pint.UnitContainer
+            Input unit.
+        registry : pint.UnitRegistry
+            the associated registry
+        **options
+            Additional options (may be ignored)
+
+        Returns
+        -------
+        out : str
+            Units following CF-Convention, using symbols.
+        """
+        import re
+
+        # convert UnitContainer back to Unit
+        unit = registry.Unit(unit)
+        # Print units using abbreviations (millimeter -> mm)
+        s = f"{unit:~D}"
+
+        # Search and replace patterns
+        pat = r"(?P<inverse>(?:1 )?/ )?(?P<unit>\w+)(?: \*\* (?P<pow>\d))?"
+
+        def repl(m):
+            i, u, p = m.groups()
+            p = p or (1 if i else "")
+            neg = "-" if i else ""
+
+            return f"{u}{neg}{p}"
+
+        out, n = re.subn(pat, repl, s)
+
+        # Remove multiplications
+        out = out.replace(" * ", " ")
+        # Delta degrees:
+        out = out.replace("Δ°", "delta_deg")
+        return out.replace("percent", "%")
+
+
+except ImportError:
+    pass
+
+
+# Reused with modification from MetPy under the terms of the BSD 3-Clause License.
+# Copyright (c) 2015,2017,2019 MetPy Developers.
 # Create registry, with preprocessors for UDUNITS-style powers (m2 s-2) and percent signs
 units = pint.UnitRegistry(
     autoconvert_offset_to_baseunit=True,
@@ -51,8 +100,7 @@
         "Import(s) unavailable to set up matplotlib support...skipping this portion "
         "of the setup."
     )
+# end of vendored code from MetPy
 
 # Set as application registry
 pint.set_application_registry(units)
-
-del pint

ci/doc.yml

@@ -17,6 +17,7 @@ dependencies:
   - ipywidgets
   - pandas
   - pooch
+  - pint
   - pip:
       - git+https://github.com/xarray-contrib/cf-xarray
       - myst_nb

doc/units.md

@@ -11,4 +11,15 @@ kernelspec:
 
 The xarray ecosystem supports unit-aware arrays using  [pint](https://pint.readthedocs.io) and [pint-xarray](https://pint-xarray.readthedocs.io). Some changes are required to make these packages work well with [UDUNITS format recommended by the CF conventions](http://cfconventions.org/Data/cf-conventions/cf-conventions-1.8/cf-conventions.html#units).
 
-`cf_xarray` makes those recommended changes when you `import cf_xarray.units`. These changes allow pint to parse UDUNIT units strings, and add several custom units like `degrees_north`, `psu` etc.
+`cf_xarray` makes those recommended changes when you `import cf_xarray.units`. These changes allow pint to parse and format UDUNIT units strings, and add several custom units like `degrees_north`, `psu` etc.
+
+## Formatting units
+
+For now, only the short format using [symbols](https://www.unidata.ucar.edu/software/udunits/udunits-2.2.28/udunits2lib.html#Syntax) is supported:
+```{code-cell}
+from pint import application_registry as ureg
+import cf_xarray.units
+
+u = ureg.Unit("m ** 3 / s ** 2")
+f"{u:cf}"
+```

doc/whats-new.rst

@@ -2,6 +2,9 @@
 
 What's New
 ----------
+v0.6.4 (unreleased)
+===================
+- integrate ``xclim``'s CF-compliant unit formatter. :pr:`278`. By `Justus Magin`_.
 
 v0.6.3 (December 16, 2021)
 ==========================