allow passing a format string to dequantify (#49)

* check that the unit formatting util works correctly

* use a format instead of casting to str

* allow passing a format to .pint.dequantify

* handle empty string formats

* move the implementation of dequantify to a shared function

* also test a empty string formatter

* actually use the format

* add a entry to whats-new.rst

Author: keewis

docs/whats-new.rst

@@ -9,6 +9,8 @@ What's new
 - allow converting indexes (except :py:class:`pandas.MultiIndex`) (:pull:`56`)
 - document the reason for requiring the ``force_ndarray_like`` or ``force_ndarray``
   options on unit registries (:pull:`59`)
+- allow passing a format string to :py:meth:`Dataset.pint.dequantify` and
+  :py:meth:`DataArray.pint.dequantify` (:pull:`49`)
 
 v0.1 (October 26 2020)
 ----------------------

pint_xarray/accessors.py

@@ -77,9 +77,11 @@ def merge_mappings(first, *mappings):
     return result
 
 
-def units_to_str_or_none(mapping):
+def units_to_str_or_none(mapping, unit_format):
+    formatter = str if not unit_format else lambda v: unit_format.format(v)
+
     return {
-        key: str(value) if isinstance(value, Unit) else value
+        key: formatter(value) if isinstance(value, Unit) else value
         for key, value in mapping.items()
     }
 
@@ -241,7 +243,7 @@ def quantify(self, units=None, unit_registry=None, **unit_kwargs):
             conversion.attach_units, units
         )
 
-    def dequantify(self):
+    def dequantify(self, format=None):
         """
         Removes units from the DataArray and its coordinates.
 
@@ -253,15 +255,19 @@ def dequantify(self):
         dequantified : DataArray
             DataArray whose array data is unitless, and of the type
             that was previously wrapped by `pint.Quantity`.
+        format : str, optional
+            The format used for the string representations.
         """
-
         units = conversion.extract_unit_attributes(self.da)
         units.update(conversion.extract_units(self.da))
 
+        unit_format = f"{{:{format}}}" if isinstance(format, str) else format
+
+        units = units_to_str_or_none(units, unit_format)
         return (
             self.da.pipe(conversion.strip_units)
             .pipe(conversion.strip_unit_attributes)
-            .pipe(conversion.attach_unit_attributes, units_to_str_or_none(units))
+            .pipe(conversion.attach_unit_attributes, units)
         )
 
     @property
@@ -504,7 +510,7 @@ def quantify(self, units=None, unit_registry=None, **unit_kwargs):
             conversion.attach_units, units
         )
 
-    def dequantify(self):
+    def dequantify(self, format=None):
         """
         Removes units from the Dataset and its coordinates.
 
@@ -516,14 +522,19 @@ def dequantify(self):
         dequantified : Dataset
             Dataset whose data variables are unitless, and of the type
             that was previously wrapped by ``pint.Quantity``.
+        format : str, optional
+            The format used for the string representations.
         """
         units = conversion.extract_unit_attributes(self.ds)
         units.update(conversion.extract_units(self.ds))
 
+        unit_format = f"{{:{format}}}" if isinstance(format, str) else format
+
+        units = units_to_str_or_none(units, unit_format)
         return (
             self.ds.pipe(conversion.strip_units)
             .pipe(conversion.strip_unit_attributes)
-            .pipe(conversion.attach_unit_attributes, units_to_str_or_none(units))
+            .pipe(conversion.attach_unit_attributes, units)
         )
 
     def to(self, units=None, **unit_kwargs):

pint_xarray/tests/test_accessors.py

@@ -6,7 +6,7 @@
 from pint.errors import UndefinedUnitError
 from xarray.testing import assert_equal
 
-from .. import conversion
+from .. import accessors, conversion
 from .utils import raises_regex
 
 pytestmark = [
@@ -101,6 +101,23 @@ def test_parse_integer_inverse(self):
         assert result.pint.units == Unit("1 / meter")
 
 
+@pytest.mark.parametrize("formatter", ("", "P", "C"))
+@pytest.mark.parametrize("flags", ("", "~", "#", "~#"))
+def test_units_to_str_or_none(formatter, flags):
+    unit_format = f"{{:{flags}{formatter}}}"
+    unit_attrs = {None: "m", "a": "s", "b": "degC", "c": "degF", "d": "degK"}
+    units = {key: unit_registry.Unit(value) for key, value in unit_attrs.items()}
+
+    expected = {key: unit_format.format(value) for key, value in units.items()}
+    actual = accessors.units_to_str_or_none(units, unit_format)
+
+    assert expected == actual
+    assert units == {key: unit_registry.Unit(value) for key, value in actual.items()}
+
+    expected = {None: None}
+    assert expected == accessors.units_to_str_or_none(expected, unit_format)
+
+
 class TestDequantifyDataArray:
     def test_strip_units(self, example_quantity_da):
         result = example_quantity_da.pint.dequantify()