Pull xarray's nbytes from nbytes attribute on arrays (#6797)

* Pull xarray's nbytes from nbytes attribute on arrays

* Calculate nbytes if it doesn't exist

* Add test

* Add docstrings

* Apply suggestions from code review

Co-authored-by: Illviljan <[email protected]>

* Add sparse variable test

* Add whats-new note

Co-authored-by: Illviljan <[email protected]>
Co-authored-by: dcherian <[email protected]>
Co-authored-by: Deepak Cherian <[email protected]>

Author: 3 people

doc/api.rst

@@ -282,7 +282,6 @@ ndarray attributes
    DataArray.shape
    DataArray.size
    DataArray.dtype
-   DataArray.nbytes
    DataArray.chunks
 
 

doc/whats-new.rst

@@ -34,6 +34,8 @@ Deprecations
 Bug fixes
 ~~~~~~~~~
 
+- :py:attr:`DataArray.nbytes` now uses the ``nbytes`` property of the underlying array if available.
+  By `Max Jones <https://github.com/maxrjones>`_.
 
 Documentation
 ~~~~~~~~~~~~~

xarray/core/dataarray.py

@@ -646,6 +646,12 @@ def size(self) -> int:
 
     @property
     def nbytes(self) -> int:
+        """
+        Total bytes consumed by the elements of this DataArray's data.
+
+        If the backend data array does not include ``nbytes``, estimates
+        the bytes consumed based on the ``size`` and ``dtype``.
+        """
         return self.variable.nbytes
 
     @property

xarray/core/dataset.py

@@ -1378,6 +1378,12 @@ def __array__(self, dtype=None):
 
     @property
     def nbytes(self) -> int:
+        """
+        Total bytes consumed by the data arrays of all variables in this dataset.
+
+        If the backend array for any variable does not include ``nbytes``, estimates
+        the total bytes for that array based on the ``size`` and ``dtype``.
+        """
         return sum(v.nbytes for v in self.variables.values())
 
     @property

xarray/core/variable.py

@@ -334,8 +334,14 @@ def shape(self):
         return self._data.shape
 
     @property
-    def nbytes(self):
-        return self.size * self.dtype.itemsize
+    def nbytes(self) -> int:
+        """
+        Total bytes consumed by the elements of the data array.
+        """
+        if hasattr(self.data, "nbytes"):
+            return self.data.nbytes
+        else:
+            return self.size * self.dtype.itemsize
 
     @property
     def _in_memory(self):

xarray/tests/test_array_api.py

@@ -43,6 +43,12 @@ def test_indexing(arrays) -> None:
     assert_equal(actual, expected)
 
 
+def test_properties(arrays) -> None:
+    np_arr, xp_arr = arrays
+    assert np_arr.nbytes == np_arr.data.nbytes
+    assert xp_arr.nbytes == np_arr.data.nbytes
+
+
 def test_reorganizing_operation(arrays) -> None:
     np_arr, xp_arr = arrays
     expected = np_arr.transpose()

xarray/tests/test_sparse.py

@@ -274,6 +274,9 @@ def setUp(self):
         self.data = sparse.random((4, 6), random_state=0, density=0.5)
         self.var = xr.Variable(("x", "y"), self.data)
 
+    def test_nbytes(self):
+        assert self.var.nbytes == self.data.nbytes
+
     def test_unary_op(self):
         assert_sparse_equal(-self.var.data, -self.data)
         assert_sparse_equal(abs(self.var).data, abs(self.data))