diff --git a/docs/whats-new.rst b/docs/whats-new.rst index b4d7ff7b..7ef7daf8 100644 --- a/docs/whats-new.rst +++ b/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) ---------------------- diff --git a/pint_xarray/accessors.py b/pint_xarray/accessors.py index 92ca2dfe..165a85bc 100644 --- a/pint_xarray/accessors.py +++ b/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): diff --git a/pint_xarray/tests/test_accessors.py b/pint_xarray/tests/test_accessors.py index e7026884..00929882 100644 --- a/pint_xarray/tests/test_accessors.py +++ b/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()