add pint to the list of supported types in SupportsArithmetic

Author: keewis

xarray/core/arithmetic.py

@@ -4,7 +4,7 @@
 import numpy as np
 
 from .options import OPTIONS
-from .pycompat import dask_array_type
+from .pycompat import dask_array_type, pint_array_type
 from .utils import not_implemented
 
 
@@ -22,22 +22,18 @@ class SupportsArithmetic:
 
     # TODO: allow extending this with some sort of registration system
     _HANDLED_TYPES = (
-        np.ndarray,
-        np.generic,
-        numbers.Number,
-        bytes,
-        str,
-    ) + dask_array_type
+        (np.ndarray, np.generic, numbers.Number, bytes, str,)
+        + dask_array_type
+        + pint_array_type
+    )
 
     def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
         from .computation import apply_ufunc
 
         # See the docstring example for numpy.lib.mixins.NDArrayOperatorsMixin.
         out = kwargs.get("out", ())
         for x in inputs + out:
-            if not isinstance(
-                x, self._HANDLED_TYPES + (SupportsArithmetic,)
-            ) and not hasattr(x, "__array_ufunc__"):
+            if not isinstance(x, self._HANDLED_TYPES + (SupportsArithmetic,)):
                 return NotImplemented
 
         if ufunc.signature is not None:

xarray/core/pycompat.py

@@ -17,3 +17,11 @@
     sparse_array_type = (sparse.SparseArray,)
 except ImportError:  # pragma: no cover
     sparse_array_type = ()
+
+try:
+    # solely for isinstance checks
+    import pint
+
+    pint_array_type = (pint.Quantity,)
+except ImportError:  # pragma: no cover
+    pint_array_type = ()