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

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

@@ -29,12 +29,14 @@
 # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 import builtins
+import operator
 from distutils.version import LooseVersion
 from typing import Union
 
-import operator
 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,9 +18,9 @@
     VectorizedIndexer,
     as_indexable,
 )
+from .npcompat import IS_NEP18_ACTIVE, _asarray
 from .options import _get_keep_attrs
 from .pycompat import dask_array_type, integer_types
-from .npcompat import IS_NEP18_ACTIVE
 from .utils import (
     OrderedSet,
     decode_numpy_dict_values,
@@ -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,160 @@
+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 create_data():
+    return (np.arange(10 * 20).reshape(10, 20) + 1) * pq.V
+
+
+def create_coord_arrays():
+    x = (np.arange(10) + 1) * pq.A
+    y = np.arange(20) + 1
+    xp = (np.arange(10) + 1) * pq.J
+    return x, y, xp
+
+
+def create_coords(x=None, y=None, xp=None):
+    x_, y_, xp_ = create_coord_arrays()
+    if x is None:
+        x = x_
+    if y is None:
+        y = y_
+    if xp is None:
+        xp = xp_
+
+    coords = dict(x=x, y=y, xp=(["x"], xp))
+    return coords
+
+
+def create_data_array(data=None, coords=None):
+    if data is None:
+        data = create_data()
+    if coords is None:
+        coords = create_coords()
+
+    return DataArray(data, dims=("x", "y"), coords=coords)
+
+
+def with_keys(mapping, keys):
+    return {key: value for key, value in mapping.items() if key in keys}
+
+
+def test_without_subclass_support():
+    with set_options(enable_experimental_ndarray_subclass_support=False):
+        data_array = create_data_array()
+        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 = create_data()
+    _, _, xp = create_coord_arrays()
+    data_array = create_data_array(data=data, coords=create_coords(xp=xp))
+
+    assert_equal_with_units(data, data_array)
+    assert_equal_with_units(xp, data_array.xp)
+
+
+def test_arithmetics():
+    v = create_data()
+    coords = create_coords()
+    da = create_data_array(data=v, coords=coords)
+
+    f = np.arange(10 * 20).reshape(10, 20) * pq.A
+    g = DataArray(f, dims=["x", "y"], coords=with_keys(coords, ["x", "y"]))
+    assert_equal_with_units(da * g, v * f)
+
+    # swapped dimension order
+    f = np.arange(20 * 10).reshape(20, 10) * pq.V
+    g = DataArray(f, dims=["y", "x"], coords=with_keys(coords, ["x", "y"]))
+    assert_equal_with_units(da + g, v + f.T)
+
+    # broadcasting
+    f = (np.arange(10) + 1) * pq.m
+    g = DataArray(f, dims=["x"], coords=with_keys(coords, ["x"]))
+    assert_equal_with_units(da / g, v / f[:, None])
+
+
+@pytest.mark.xfail(reason="units don't survive through combining yet")
+def test_combine():
+    from xarray import concat
+
+    data_array = create_data_array()
+
+    a = data_array[:, :10]
+    b = data_array[:, 10:]
+
+    assert_equal_with_units(concat([a, b], dim="y"), data_array)
+
+
+def test_unit_checking():
+    coords = create_coords()
+    da = create_data_array(coords=coords)
+
+    f = np.arange(10 * 20).reshape(10, 20) * pq.A
+    g = DataArray(f, dims=["x", "y"], coords=with_keys(coords, ["x", "y"]))
+    with pytest.raises(ValueError, match="Unable to convert between units"):
+        da + g
+
+
+@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.
+    """
+    x, *_ = create_coord_arrays()
+    data_array = create_data_array(coords=create_coords(x=x))
+    assert_equal_with_units(data_array.x, x)
+
+
+def test_sel():
+    data = create_data()
+    _, y, _ = create_coord_arrays()
+    data_array = create_data_array(data=data, coords=create_coords(y=y))
+    assert_equal_with_units(data_array.sel(y=y[0]), data[:, 0])
+
+
+def test_mean():
+    data = create_data()
+    data_array = create_data_array(data=data)
+    assert_equal_with_units(data_array.mean("x"), data.mean(0))