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

Fix the units and dimensionality properties #31

Merged
merged 12 commits into from
Oct 4, 2020
5 changes: 4 additions & 1 deletion docs/whats-new.rst
Original file line number Diff line number Diff line change
Expand Up @@ -9,4 +9,7 @@ What's new
- rewrite :py:meth:`DataArray.pint.quantify`,
:py:meth:`Dataset.pint.quantify`, :py:meth:`DataArray.pint.dequantify` and
:py:meth:`Dataset.pint.dequantify` (:pull:`17`)
- expose :py:func:`pint_xarray.testing.assert_units_equal` as public API (:pull:`24`)
- expose :py:func:`pint_xarray.testing.assert_units_equal` as public API (:pull:`24`)
- fix the :py:attr:`DataArray.pint.units`, :py:attr:`DataArray.pint.magnitude`
and :py:attr:`DataArray.pint.dimensionality` properties and add docstrings for
all three. (:pull:`31`)
20 changes: 12 additions & 8 deletions pint_xarray/accessors.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
import pint
from pint.quantity import Quantity
from pint.unit import Unit
from xarray import DataArray, register_dataarray_accessor, register_dataset_accessor
from xarray import register_dataarray_accessor, register_dataset_accessor

from . import conversion

Expand Down Expand Up @@ -230,22 +230,26 @@ def dequantify(self):

@property
def magnitude(self):
return self.da.data.magnitude
"""the magnitude of the data or the data itself if not a quantity."""
data = self.da.data
return getattr(data, "magnitude", data)

@property
def units(self):
return self.da.data.units
"""the units of the data or :py:obj:`None` if not a quantity.

Setting the units is possible, but only if the data is not already a quantity.
"""
return getattr(self.da.data, "units", None)

@units.setter
def units(self, units):
quantity = conversion.array_attach_units(self.da.data, units)
self.da = DataArray(
dim=self.da.dims, data=quantity, coords=self.da.coords, attrs=self.da.attrs
)
self.da.data = conversion.array_attach_units(self.da.data, units)

@property
def dimensionality(self):
return self.da.data.dimensionality
"""get the dimensionality of the data or :py:obj:`None` if not a quantity."""
return getattr(self.da.data, "dimensionality", None)

@property
def registry(self):
Expand Down
31 changes: 31 additions & 0 deletions pint_xarray/tests/test_accessors.py
Original file line number Diff line number Diff line change
Expand Up @@ -123,6 +123,37 @@ def test_roundtrip_data(self, example_unitless_da):
assert_equal(result, orig)


class TestPropertiesDataArray:
def test_magnitude_getattr(self, example_quantity_da):
da = example_quantity_da
actual = da.pint.magnitude
assert not isinstance(actual, Quantity)

def test_magnitude_getattr_unitless(self, example_unitless_da):
da = example_unitless_da
xr.testing.assert_duckarray_equal(da.pint.magnitude, da.data)

def test_units_getattr(self, example_quantity_da):
da = example_quantity_da
actual = da.pint.units
assert isinstance(actual, Unit)
assert actual == unit_registry.m

def test_units_setattr(self, example_quantity_da):
da = example_quantity_da
with pytest.raises(ValueError):
da.pint.units = "s"

def test_units_getattr_unitless(self, example_unitless_da):
da = example_unitless_da
assert da.pint.units is None

def test_units_setattr_unitless(self, example_unitless_da):
da = example_unitless_da
da.pint.units = unit_registry.s
assert da.pint.units == unit_registry.s


@pytest.fixture()
def example_unitless_ds():
users = np.linspace(0, 10, 20)
Expand Down