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]>
pydata · Jul 22, 2022 · 60f8c3d · 60f8c3d
1 parent ed56df2
commit 60f8c3d
Show file tree

Hide file tree

Showing 7 changed files with 31 additions and 3 deletions.
diff --git a/doc/api.rst b/doc/api.rst
@@ -282,7 +282,6 @@ ndarray attributes
    DataArray.shape
    DataArray.size
    DataArray.dtype
-   DataArray.nbytes
    DataArray.chunks
 
 

diff --git a/doc/whats-new.rst b/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
 ~~~~~~~~~~~~~

diff --git a/xarray/core/dataarray.py b/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

diff --git a/xarray/core/dataset.py b/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

diff --git a/xarray/core/variable.py b/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):

diff --git a/xarray/tests/test_array_api.py b/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()

diff --git a/xarray/tests/test_sparse.py b/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))