Picking up #1118: Do not convert subclasses of ndarray unless required

ci/requirements/py36.yml

@@ -35,4 +35,4 @@ dependencies:
   - iris>=1.10
   - pydap
   - lxml
-
+  - quantities

ci/requirements/py37.yml

@@ -32,5 +32,6 @@ dependencies:
   - cfgrib>=0.9.2
   - lxml
   - pydap
+  - quantities
   - pip:
     - numbagg

setup.cfg

@@ -82,6 +82,8 @@ ignore_missing_imports = True
 ignore_missing_imports = True
 [mypy-pytest.*]
 ignore_missing_imports = True
+[mypy-quantities.*]
+ignore_missing_imports = True
 [mypy-rasterio.*]
 ignore_missing_imports = True
 [mypy-scipy.*]

xarray/core/duck_array_ops.py

@@ -162,10 +162,12 @@ def trapz(y, x, axis):
 
 
 def asarray(data):
+    from .npcompat import _asarray
+
     return (
         data
         if (isinstance(data, dask_array_type) or hasattr(data, "__array_function__"))
-        else np.asarray(data)
+        else _asarray(data)
     )
 
 

xarray/core/npcompat.py

@@ -35,6 +35,8 @@
 
 import numpy as np
 
+from .options import OPTIONS
+
 try:
     from numpy import isin
 except ImportError:
@@ -378,3 +380,12 @@ def __array_function__(self, *args, **kwargs):
 
 
 IS_NEP18_ACTIVE = _is_nep18_active()
+
+
+def _asarray(data):
+    # options get set after import, so this needs to be done in a
+    # function
+    if OPTIONS["enable_experimental_ndarray_subclass_support"]:
+        return np.asanyarray(data)
+    else:
+        return np.asarray(data)

xarray/core/options.py

@@ -8,6 +8,9 @@
 CMAP_SEQUENTIAL = "cmap_sequential"
 CMAP_DIVERGENT = "cmap_divergent"
 KEEP_ATTRS = "keep_attrs"
+ENABLE_EXPERIMENTAL_NDARRAY_SUBCLASS_SUPPORT = (
+    "enable_experimental_ndarray_subclass_support"
+)
 
 
 OPTIONS = {
@@ -19,6 +22,7 @@
     CMAP_SEQUENTIAL: "viridis",
     CMAP_DIVERGENT: "RdBu_r",
     KEEP_ATTRS: "default",
+    ENABLE_EXPERIMENTAL_NDARRAY_SUBCLASS_SUPPORT: False,
 }
 
 _JOIN_OPTIONS = frozenset(["inner", "outer", "left", "right", "exact"])
@@ -35,6 +39,7 @@ def _positive_integer(value):
     FILE_CACHE_MAXSIZE: _positive_integer,
     WARN_FOR_UNCLOSED_FILES: lambda value: isinstance(value, bool),
     KEEP_ATTRS: lambda choice: choice in [True, False, "default"],
+    ENABLE_EXPERIMENTAL_NDARRAY_SUBCLASS_SUPPORT: lambda value: isinstance(value, bool),
 }
 
 
@@ -98,6 +103,9 @@ class set_options:
       attrs, ``False`` to always discard them, or ``'default'`` to use original
       logic that attrs should only be kept in unambiguous circumstances.
       Default: ``'default'``.
+    - ``enable_experimental_ndarray_subclass_support``: whether or not
+      to enable the support for subclasses of numpy's ndarray.
+      Default: ``False``.
 
     You can use ``set_options`` either as a context manager:
 

xarray/core/variable.py

@@ -18,7 +18,7 @@
     VectorizedIndexer,
     as_indexable,
 )
-from .npcompat import IS_NEP18_ACTIVE
+from .npcompat import IS_NEP18_ACTIVE, _asarray
 from .options import _get_keep_attrs
 from .pycompat import dask_array_type, integer_types
 from .utils import (
@@ -193,6 +193,8 @@ def as_compatible_data(data, fastpath=False):
             data[mask] = fill_value
         else:
             data = np.asarray(data)
+    elif isinstance(data, np.matrix):
+        data = np.asarray(data)
 
     if not isinstance(data, np.ndarray):
         if hasattr(data, "__array_function__"):
@@ -208,8 +210,7 @@ def as_compatible_data(data, fastpath=False):
                 )
 
     # validate whether the data is valid data types
-    data = np.asarray(data)
-
+    data = _asarray(data)
     if isinstance(data, np.ndarray):
         if data.dtype.kind == "O":
             data = _possibly_convert_objects(data)
@@ -244,6 +245,21 @@ def _as_array_or_item(data):
     return data
 
 
+def _as_any_array_or_item(data):
+    """ Return the given values as a numpy array subclass instance, or
+    individual item if it's a 0d datetime64 or timedelta64 array.
+
+    The same caveats as for ``_as_array_or_item`` apply.
+    """
+    data = _asarray(data)
+    if data.ndim == 0:
+        if data.dtype.kind == "M":
+            data = np.datetime64(data, "ns")
+        elif data.dtype.kind == "m":
+            data = np.timedelta64(data, "ns")
+    return data
+
+
 class Variable(
     common.AbstractArray, arithmetic.SupportsArithmetic, utils.NdimSizeLenMixin
 ):
@@ -322,7 +338,7 @@ def data(self):
         ):
             return self._data
         else:
-            return self.values
+            return _as_any_array_or_item(self._data)
 
     @data.setter
     def data(self, data):

xarray/tests/test_quantities.py

@@ -0,0 +1,140 @@
+import numpy as np
+import pytest
+
+from xarray import DataArray, set_options
+
+try:
+    import quantities as pq
+
+    has_quantities = True
+except ImportError:
+    has_quantities = False
+
+pytestmark = pytest.mark.skipif(not has_quantities, reason="requires python-quantities")
+
+
+set_options(enable_experimental_ndarray_subclass_support=True)
+
+
+def assert_equal_with_units(a, b):
+    a = a if not isinstance(a, DataArray) else a.data
+    b = b if not isinstance(b, DataArray) else b.data
+
+    assert (hasattr(a, "units") and hasattr(b, "units")) and a.units == b.units
+
+    assert (hasattr(a, "magnitude") and hasattr(b, "magnitude")) and np.allclose(
+        a.magnitude, b.magnitude
+    )
+
+
+def test_without_subclass_support():
+    with set_options(enable_experimental_ndarray_subclass_support=False):
+        data_array = DataArray(data=np.arange(10) * pq.m)
+        assert not hasattr(data_array.data, "units")
+
+
+@pytest.mark.filterwarnings("ignore:the matrix subclass:PendingDeprecationWarning")
+def test_matrix():
+    matrix = np.matrix([[1, 2], [3, 4]])
+    da = DataArray(matrix)
+
+    assert not isinstance(da.data, np.matrix)
+
+
+def test_masked_array():
+    masked = np.ma.array([[1, 2], [3, 4]], mask=[[0, 1], [1, 0]])
+    da = DataArray(masked)
+    assert not isinstance(da.data, np.ma.MaskedArray)
+
+
+def test_units_in_data_and_coords():
+    data = np.arange(10) * pq.m
+    x = np.arange(10) * pq.s
+    xp = x.rescale(pq.ms)
+    data_array = DataArray(data=data, coords={"x": x, "xp": ("x", xp)}, dims=["x"])
+
+    assert_equal_with_units(data, data_array)
+    assert_equal_with_units(xp, data_array.xp)
+
+
+def test_arithmetics():
+    x = np.arange(10)
+    y = np.arange(20)
+
+    coords = {"x": x, "y": y}
+
+    f = (np.arange(10 * 20).reshape(10, 20) + 1) * pq.V
+    g = np.arange(10 * 20).reshape(10, 20) * pq.A
+
+    a = DataArray(data=f, coords=coords, dims=("x", "y"))
+    b = DataArray(data=g, coords=coords, dims=("x", "y"))
+
+    assert_equal_with_units(a * b, f * g)
+
+    # swapped dimension order
+    g = np.arange(20 * 10).reshape(20, 10) * pq.V
+    b = DataArray(data=g, coords=coords, dims=("y", "x"))
+    assert_equal_with_units(a + b, f + g.T)
+
+    # broadcasting
+    g = (np.arange(10) + 1) * pq.m
+    b = DataArray(data=g, coords={"x": x}, dims=["x"])
+    assert_equal_with_units(a / b, f / g[:, None])
+
+
+@pytest.mark.xfail(reason="units don't survive through combining yet")
+def test_combine():
+    from xarray import concat
+
+    data = (np.arange(15) + 10) * pq.m
+    y = np.arange(len(data))
+
+    data_array = DataArray(data=data, coords={"y": y}, dims=["y"])
+    a = data_array[:5]
+    b = data_array[5:]
+
+    assert_equal_with_units(concat([a, b], dim="y"), data_array)
+
+
+def test_unit_checking():
+    coords = {"x": np.arange(10), "y": np.arange(20)}
+
+    f = np.arange(10 * 20).reshape(10, 20) * pq.A
+    g = np.arange(10 * 20).reshape(10, 20) * pq.V
+
+    a = DataArray(f, coords=coords, dims=("x", "y"))
+    b = DataArray(g, coords=coords, dims=("x", "y"))
+    with pytest.raises(ValueError, match="Unable to convert between units"):
+        a + b
+
+
+@pytest.mark.xfail(reason="units in indexes not supported")
+def test_units_in_indexes():
+    """ Test if units survive through xarray indexes.
+    Indexes are borrowed from Pandas, and Pandas does not support
+    units. Therefore, we currently don't intend to support units on
+    indexes either.
+    """
+    data = np.arange(15) * 10
+    x = np.arange(len(data)) * pq.A
+
+    data_array = DataArray(data=data, coords={"x": x}, dims=["x"])
+    assert_equal_with_units(data_array.x, x)
+
+
+def test_sel():
+    data = np.arange(10 * 20).reshape(10, 20) * pq.m / pq.s
+    x = np.arange(data.shape[0]) * pq.m
+    y = np.arange(data.shape[1]) * pq.s
+
+    data_array = DataArray(data=data, coords={"x": x, "y": y}, dims=("x", "y"))
+    assert_equal_with_units(data_array.sel(y=y[0]), data[:, 0])
+
+
+def test_mean():
+    data = np.arange(10 * 20).reshape(10, 20) * pq.V
+    x = np.arange(data.shape[0])
+    y = np.arange(data.shape[1])
+
+    data_array = DataArray(data=data, coords={"x": x, "y": y}, dims=("x", "y"))
+    assert_equal_with_units(data_array.mean("x"), data.mean(0))