Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

allow passing a format string to dequantify #49

Merged
merged 11 commits into from
Feb 8, 2021
2 changes: 2 additions & 0 deletions docs/whats-new.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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)
----------------------
Expand Down
25 changes: 18 additions & 7 deletions pint_xarray/accessors.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
}

Expand Down Expand Up @@ -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.

Expand All @@ -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
Expand Down Expand Up @@ -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.

Expand All @@ -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):
Expand Down
19 changes: 18 additions & 1 deletion pint_xarray/tests/test_accessors.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = [
Expand Down Expand Up @@ -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()
Expand Down