Repository: hgrecco/pint

Issue:
Preserve `ndarray`-like magnitudes or reject them from the start
I am trying to use a [`xarray.DataArray`](http://xarray.pydata.org/en/latest/api.html#dataarray) magnitude with a `pint` quantity.  It appears to work, but doesn't really, because [`Quantity.__getattr__` forces it to convert to `numpy.ndarray`](https://github.com/hgrecco/pint/blob/master/pint/quantity.py#L1167).  This leads to conversion at very unexpected behaviour:

```
    In [526]: q = ureg.Quantity(xarray.DataArray([1, 2, 3]), 'm')
    
    In [527]: print(q)
    <xarray.DataArray (dim_0: 3)>array([1, 2, 3])Dimensions without coordinates: dim_0 meter
    
    In [528]: a = array(q)
    
    In [529]: print(q)
    [1 2 3] meter
```

I think pint should either accept magnitudes that behave like `numpy.ndarray` but don't subclass it (such as `xarray.DataArray`), or reject them when the object is constructed.  I don't see how converting it in the middle of attribute lookup can do any good…

Discussion:

Patch:
diff --git a/.travis.yml b/.travis.yml
index 4f8613f..a6cf799 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -29,6 +29,7 @@ env:
   - PKGS="python=3.6 numpy uncertainties"
   - PKGS="python=3.7 numpy uncertainties"
   - PKGS="python=3.8 numpy uncertainties"
+  - PKGS="python xarray netCDF4"
 
   # TODO: pandas tests
   # - PKGS="python=3.7 numpy pandas uncertainties pandas"
diff --git a/CHANGES b/CHANGES
index 20689cd..c131d45 100644
--- a/CHANGES
+++ b/CHANGES
@@ -4,9 +4,16 @@ Pint Changelog
 0.10 (unreleased)
 -----------------
 
+- Improved compatbility for upcast types like xarray's DataArray or Dataset, to which
+  Pint Quantities now fully defer for arithmetic and NumPy operations. A collection of
+  basic tests for proper deferral has been added (for full integration tests, see
+  xarray's test suite). The list of upcast types is available at
+  `pint.compat.upcast_types` in the API.
+  (Issue #959, Thanks Jon Thielen)
 - Moved docstrings to Numpy Docs
+  (Issue #958)
 - Added tests for immutability of the magnitude's type under common operations
-  (Issue #956, Thanks Jon Thielen)
+  (Issue #957, Thanks Jon Thielen)
 - Switched test configuration to pytest and added tests of Pint's matplotlib support.
   (Issue #954, Thanks Jon Thielen)
 - Deprecate array protocol fallback except where explicitly defined (`__array__`,
diff --git a/docs/numpy.rst b/docs/numpy.rst
index f849af2..c9a46b9 100644
--- a/docs/numpy.rst
+++ b/docs/numpy.rst
@@ -144,7 +144,11 @@ upcast types to which Pint defers (see
 
 - ``PintArray``, as defined by pint-pandas
 - ``Series``, as defined by pandas
-- ``DataArray``, as defined by xarray
+- ``DataArray``, ``Dataset``, and ``Variable``, as defined by xarray
+
+If your application requires extension of this collection of types, it is available in
+Pint's API at ``pint.compat.upcast_types``. Note that these are also the types to which
+a Quantity object will defer for arithmetic operations.
 
 To achive these function and ufunc overrides, Pint uses the ``__array_function__`` and
 ``__array_ufunc__`` protocols respectively, as recommened by NumPy. This means that
diff --git a/pint/compat.py b/pint/compat.py
index 7642b15..e87c26c 100644
--- a/pint/compat.py
+++ b/pint/compat.py
@@ -149,11 +149,36 @@ if not HAS_BABEL:
     babel_parse = babel_units = missing_dependency("Babel")  # noqa: F811
 
 # Define location of pint.Quantity in NEP-13 type cast hierarchy by defining upcast and
-# downcast/wrappable types
+# downcast/wrappable types using guarded imports
+upcast_types = []
+
+# pint-pandas (PintArray)
+try:
+    from pintpandas import PintArray
+
+    upcast_types.append(PintArray)
+except ImportError:
+    pass
+
+# Pandas (Series)
+try:
+    from pandas import Series
+
+    upcast_types.append(Series)
+except ImportError:
+    pass
+
+# xarray (DataArray, Dataset, Variable)
+try:
+    from xarray import DataArray, Dataset, Variable
+
+    upcast_types += [DataArray, Dataset, Variable]
+except ImportError:
+    pass
 
 
 def is_upcast_type(other):
-    """Check if the type object is a upcast type.
+    """Check if the type object is a upcast type using preset list.
 
     Parameters
     ----------
@@ -163,8 +188,7 @@ def is_upcast_type(other):
     -------
     bool
     """
-    # Check if class name is in preset list
-    return other.__name__ in ("PintArray", "Series", "DataArray")
+    return other in upcast_types
 
 
 def eq(lhs, rhs, check_all):
diff --git a/pint/quantity.py b/pint/quantity.py
index 8bdf737..1d6a372 100644
--- a/pint/quantity.py
+++ b/pint/quantity.py
@@ -158,7 +158,9 @@ class Quantity(PrettyIPython, SharedRegistryObject):
     def __new__(cls, value, units=None):
         global SKIP_ARRAY_FUNCTION_CHANGE_WARNING
 
-        if units is None:
+        if is_upcast_type(type(value)):
+            raise TypeError(f"Quantity cannot wrap upcast type {type(value)}")
+        elif units is None:
             if isinstance(value, str):
                 if value == "":
                     raise ValueError(
@@ -712,6 +714,7 @@ class Quantity(PrettyIPython, SharedRegistryObject):
             return complex(self._convert_magnitude_not_inplace(UnitsContainer()))
         raise DimensionalityError(self._units, "dimensionless")
 
+    @check_implemented
     def _iadd_sub(self, other, op):
         """Perform addition or subtraction operation in-place and return the result.
 
@@ -964,6 +967,7 @@ class Quantity(PrettyIPython, SharedRegistryObject):
         else:
             return -self._add_sub(other, operator.sub)
 
+    @check_implemented
     @ireduce_dimensions
     def _imul_div(self, other, magnitude_op, units_op=None):
         """Perform multiplication or division operation in-place and return the
@@ -1183,6 +1187,7 @@ class Quantity(PrettyIPython, SharedRegistryObject):
             raise DimensionalityError(self._units, "dimensionless")
         return self.__class__(magnitude, UnitsContainer({}))
 
+    @check_implemented
     def __imod__(self, other):
         if not self._check(other):
             other = self.__class__(other, UnitsContainer({}))
@@ -1226,6 +1231,7 @@ class Quantity(PrettyIPython, SharedRegistryObject):
             raise DimensionalityError(self._units, "dimensionless")
         return (self.__class__(q, UnitsContainer({})), self.__class__(r, unit))
 
+    @check_implemented
     def __ipow__(self, other):
         if not isinstance(self._magnitude, ndarray):
             return self.__pow__(other)
@@ -1406,6 +1412,7 @@ class Quantity(PrettyIPython, SharedRegistryObject):
         except DimensionalityError:
             return False
 
+    @check_implemented
     def __ne__(self, other):
         out = self.__eq__(other)
         if isinstance(out, ndarray):


Test patch:
diff --git a/pint/testsuite/test_compat_upcast.py b/pint/testsuite/test_compat_upcast.py
new file mode 100644
index 0000000..7e3cd10
--- /dev/null
+++ b/pint/testsuite/test_compat_upcast.py
@@ -0,0 +1,127 @@
+import pytest
+
+from pint import UnitRegistry
+
+# Conditionally import NumPy and any upcast type libraries
+np = pytest.importorskip("numpy", reason="NumPy is not available")
+xr = pytest.importorskip("xarray", reason="xarray is not available")
+
+# Set up unit registry and sample
+ureg = UnitRegistry()
+q = [[1.0, 2.0], [3.0, 4.0]] * ureg.m
+
+
+@pytest.fixture
+def da():
+    return xr.DataArray(q.copy())
+
+
+@pytest.fixture
+def ds():
+    return xr.tutorial.load_dataset("air_temperature")
+
+
+def test_xarray_quantity_creation():
+    with pytest.raises(TypeError) as exc:
+        ureg.Quantity(xr.DataArray(np.arange(4)), "m")
+        assert "Quantity cannot wrap upcast type" in str(exc)
+    assert xr.DataArray(q).data is q
+
+
+def test_quantification(ds):
+    da = ds["air"][0]
+    da.data = ureg.Quantity(da.values, da.attrs.pop("units"))
+    mean = da.mean().item()
+    assert mean.units == ureg.K
+    assert np.isclose(mean, 274.166259765625 * ureg.K)
+
+
+@pytest.mark.parametrize(
+    "op",
+    [
+        lambda x, y: x + y,
+        lambda x, y: x - (-y),
+        lambda x, y: x * y,
+        lambda x, y: x / (y ** -1),
+    ],
+)
+@pytest.mark.parametrize(
+    "pair",
+    [
+        (q, xr.DataArray(q)),
+        (
+            xr.DataArray([1.0, 2.0] * ureg.m, dims=("y",)),
+            xr.DataArray(
+                np.arange(6, dtype="float").reshape(3, 2, 1), dims=("z", "y", "x")
+            )
+            * ureg.km,
+        ),
+        (1 * ureg.m, xr.DataArray(q)),
+    ],
+)
+def test_binary_arithmetic_commutativity(op, pair):
+    z0 = op(*pair)
+    z1 = op(*pair[::-1])
+    z1 = z1.transpose(*z0.dims)
+    assert np.all(np.isclose(z0.data, z1.data.to(z0.data.units)))
+
+
+def test_eq_commutativity(da):
+    assert np.all((q.T == da) == (da.transpose() == q))
+
+
+def test_ne_commutativity(da):
+    assert np.all((q != da.transpose()) == (da != q.T))
+
+
+def test_dataset_operation_with_unit(ds):
+    ds0 = ureg.K * ds.isel(time=0)
+    ds1 = (ds * ureg.K).isel(time=0)
+    xr.testing.assert_identical(ds0, ds1)
+    assert np.isclose(ds0["air"].mean().item(), 274.166259765625 * ureg.K)
+
+
+def test_dataarray_inplace_arithmetic_roundtrip(da):
+    da_original = da.copy()
+    q_to_modify = q.copy()
+    da += q
+    xr.testing.assert_identical(da, xr.DataArray([[2, 4], [6, 8]] * ureg.m))
+    da -= q
+    xr.testing.assert_identical(da, da_original)
+    da *= ureg.m
+    xr.testing.assert_identical(da, xr.DataArray(q * ureg.m))
+    da /= ureg.m
+    xr.testing.assert_identical(da, da_original)
+    # Operating inplace with DataArray converts to DataArray
+    q_to_modify += da
+    q_to_modify -= da
+    assert np.all(np.isclose(q_to_modify.data, q))
+
+
+def test_dataarray_inequalities(da):
+    xr.testing.assert_identical(
+        2 * ureg.m > da, xr.DataArray([[True, False], [False, False]])
+    )
+    xr.testing.assert_identical(
+        2 * ureg.m < da, xr.DataArray([[False, False], [True, True]])
+    )
+    with pytest.raises(ValueError) as exc:
+        da > 2
+        assert "Cannot compare Quantity and <class 'int'>" in str(exc)
+
+
+def test_array_function_deferral(da):
+    lower = 2 * ureg.m
+    upper = 3 * ureg.m
+    args = (da, lower, upper)
+    assert (
+        lower.__array_function__(
+            np.clip, tuple(set(type(arg) for arg in args)), args, {}
+        )
+        is NotImplemented
+    )
+
+
+def test_array_ufunc_deferral(da):
+    lower = 2 * ureg.m
+    assert lower.__array_ufunc__(np.maximum, "__call__", lower, da) is NotImplemented
diff --git a/pint/testsuite/test_quantity.py b/pint/testsuite/test_quantity.py
index d820964..9d61bbe 100644
--- a/pint/testsuite/test_quantity.py
+++ b/pint/testsuite/test_quantity.py
@@ -12,6 +12,12 @@ from pint.testsuite.parameterized import ParameterizedTestCase
 from pint.unit import UnitsContainer
 
 
+class FakeWrapper:
+    # Used in test_upcast_type_rejection_on_creation
+    def __init__(self, q):
+        self.q = q
+
+
 class TestQuantity(QuantityTestCase):
 
     FORCE_NDARRAY = False
@@ -536,6 +542,11 @@ class TestQuantity(QuantityTestCase):
     def test_no_ndarray_coercion_without_numpy(self):
         self.assertRaises(ValueError, self.Q_(1, "m").__array__)
 
+    @patch("pint.compat.upcast_types", [FakeWrapper])
+    def test_upcast_type_rejection_on_creation(self):
+        self.assertRaises(TypeError, self.Q_, FakeWrapper(42), "m")
+        self.assertEqual(FakeWrapper(self.Q_(42, "m")).q, self.Q_(42, "m"))
+
 
 class TestQuantityToCompact(QuantityTestCase):
     def assertQuantityAlmostIdentical(self, q1, q2):