diff --git a/docs/whats-new.rst b/docs/whats-new.rst index 4863d6a6..8c219a08 100644 --- a/docs/whats-new.rst +++ b/docs/whats-new.rst @@ -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`) \ No newline at end of file +- 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`) diff --git a/pint_xarray/accessors.py b/pint_xarray/accessors.py index ad623ca4..fa97aac7 100644 --- a/pint_xarray/accessors.py +++ b/pint_xarray/accessors.py @@ -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 @@ -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): diff --git a/pint_xarray/tests/test_accessors.py b/pint_xarray/tests/test_accessors.py index 8867af8a..685b0e34 100644 --- a/pint_xarray/tests/test_accessors.py +++ b/pint_xarray/tests/test_accessors.py @@ -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)