first tests

Author: TomNicholas

pintxarray/tests/__init__.py

@@ -0,0 +1,112 @@
+import pytest
+
+import xarray as xr
+from xarray.testing import assert_equal
+
+import numpy as np
+from numpy.testing import assert_array_equal
+
+from pint import UnitRegistry
+from pint.unit import Unit
+from pint.errors import DimensionalityError
+
+from pintxarray.accessors import PintDataArrayAccessor, PintDatasetAccessor
+from .utils import extract_units
+
+
+# make sure scalars are converted to 0d arrays so quantities can
+# always be treated like ndarrays
+unit_registry = UnitRegistry(force_ndarray=True)
+Quantity = unit_registry.Quantity
+
+
+@pytest.fixture()
+def example_unitless_da():
+    array = np.linspace(0, 10, 20)
+    x = np.arange(20)
+    da = xr.DataArray(data=array, dims="x", coords={"x": x})
+    da.attrs['units'] = 'm'
+    da.coords['x'].attrs['units'] = 's'
+    return da
+
+@pytest.fixture()
+def example_quantity_da():
+    array = np.linspace(0, 10, 20) * unit_registry.m
+    x = np.arange(20) * unit_registry.s
+    return xr.DataArray(data=array, dims="x", coords={"x": x})
+
+
+class TestQuantifyDataArray:
+    @pytest.mark.skip(reason="Not yet implemented")
+    def test_attach_units_given_registry(self):
+        ...
+
+    def test_attach_units(self, example_unitless_da):
+        orig = example_unitless_da
+        result = orig.pint.quantify('m')
+        assert_array_equal(result.data.magnitude, orig.data)
+        assert str(result.data.units) == 'meter'
+
+    def test_attach_units_from_attrs(self):
+        ...
+
+    def test_attach_unit_class(self):
+        ...
+
+    def test_error_when_already_units(self, example_quantity_da):
+        da = example_quantity_da
+        with pytest.raises(ValueError):
+            da.pint.quantify()
+
+    def test_error_on_nonsense_units(self):
+        ...
+
+    def test_registry_kwargs(self):
+        ...
+
+
+@pytest.mark.skip(reason="Not yet implemented")
+class TestDequantifyDataArray:
+    def test_strip_units(self):
+        ...
+
+    def test_error_if_no_units(self):
+        ...
+
+    def test_roundtrip_attrs(self):
+        ...
+
+
+@pytest.mark.skip(reason="Not yet implemented")
+class TestPropertiesDataArray:
+    def test_units(self):
+        ...
+
+
+@pytest.mark.skip(reason="Not yet implemented")
+class TestConversionDataArray:
+    def test_units(self):
+        ...
+
+
+@pytest.mark.skip(reason="Not yet implemented")
+class TestUncertainties:
+    def test_units(self):
+        ...
+
+
+@pytest.mark.skip(reason="Not yet implemented")
+class TestQuantifyDataSet:
+    def test_attach_units(self):
+        ...
+
+    def test_attach_str_units(self):
+        ...
+
+    def test_attach_units_from_registry(self):
+        ...
+
+
+@pytest.mark.skip(reason="Not yet implemented")
+class TestIndexing:
+    ...

pintxarray/tests/test_accessors.py

@@ -0,0 +1,80 @@
+import xarray as xr
+
+from pint.quantity import Quantity
+
+
+def array_extract_units(obj):
+    if isinstance(obj, (xr.Variable, xr.DataArray, xr.Dataset)):
+        obj = obj.data
+
+    try:
+        return obj.units
+    except AttributeError:
+        return None
+
+def extract_units(obj):
+    if isinstance(obj, xr.Dataset):
+        vars_units = {
+            name: array_extract_units(value) for name, value in obj.data_vars.items()
+        }
+        coords_units = {
+            name: array_extract_units(value) for name, value in obj.coords.items()
+        }
+
+        units = {**vars_units, **coords_units}
+    elif isinstance(obj, xr.DataArray):
+        vars_units = {obj.name: array_extract_units(obj)}
+        coords_units = {
+            name: array_extract_units(value) for name, value in obj.coords.items()
+        }
+
+        units = {**vars_units, **coords_units}
+    elif isinstance(obj, xr.Variable):
+        vars_units = {None: array_extract_units(obj.data)}
+
+        units = {**vars_units}
+    elif isinstance(obj, Quantity):
+        vars_units = {None: array_extract_units(obj)}
+
+        units = {**vars_units}
+    else:
+        units = {}
+
+    return units
+
+def assert_units_equal(a, b):
+    __tracebackhide__ = True
+    assert extract_units(a) == extract_units(b)
+
+def assert_equal_with_units(a, b):
+    # works like xr.testing.assert_equal, but also explicitly checks units
+    # so, it is more like assert_identical
+    __tracebackhide__ = True
+
+    if isinstance(a, xr.Dataset) or isinstance(b, xr.Dataset):
+        a_units = extract_units(a)
+        b_units = extract_units(b)
+
+        a_without_units = strip_units(a)
+        b_without_units = strip_units(b)
+
+        assert a_without_units.equals(b_without_units), formatting.diff_dataset_repr(
+            a, b, "equals"
+        )
+        assert a_units == b_units
+    else:
+        a = a if not isinstance(a, (xr.DataArray, xr.Variable)) else a.data
+        b = b if not isinstance(b, (xr.DataArray, xr.Variable)) else b.data
+
+        assert type(a) == type(b) or (
+            isinstance(a, Quantity) and isinstance(b, Quantity)
+        )
+
+        # workaround until pint implements allclose in __array_function__
+        if isinstance(a, Quantity) or isinstance(b, Quantity):
+            assert (
+                hasattr(a, "magnitude") and hasattr(b, "magnitude")
+            ) and np.allclose(a.magnitude, b.magnitude, equal_nan=True)
+            assert (hasattr(a, "units") and hasattr(b, "units")) and a.units == b.units
+        else:
+            assert np.allclose(a, b, equal_nan=True)