Recreate @gajomi's pydata#2070 to keep attrs when calling astype()

Author: dnowacki-usgs

xarray/core/common.py

@@ -1305,6 +1305,46 @@ def isin(self, test_elements):
             dask="allowed",
         )
 
+    def astype(self, dtype, casting="unsafe", copy=True):
+        """
+        Copy of the xarray object, with data cast to a specified type.
+        Leaves coordinate dtype unchanged.
+
+        Parameters
+        ----------
+        dtype : str or dtype
+            Typecode or data-type to which the array is cast.
+        casting : {'no', 'equiv', 'safe', 'same_kind', 'unsafe'}, optional
+            Controls what kind of data casting may occur. Defaults to 'unsafe'
+            for backwards compatibility.
+
+            * 'no' means the data types should not be cast at all.
+            * 'equiv' means only byte-order changes are allowed.
+            * 'safe' means only casts which can preserve values are allowed.
+            * 'same_kind' means only safe casts or casts within a kind,
+                like float64 to float32, are allowed.
+            * 'unsafe' means any data conversions may be done.
+        copy : bool, optional
+            By default, astype always returns a newly allocated array. If this
+            is set to False and the `dtype` requirement is satisfied, the input
+            array is returned instead of a copy.
+
+        See also
+        --------
+        np.ndarray.astype
+        dask.array.Array.astype
+        """
+        from .computation import apply_ufunc
+
+        return apply_ufunc(
+            duck_array_ops.astype,
+            self,
+            dtype,
+            keep_attrs=True,
+            kwargs={"casting": casting, "copy": copy},
+            dask="allowed",
+        )
+
     def __enter__(self: T) -> T:
         return self
 

xarray/core/duck_array_ops.py

@@ -149,6 +149,10 @@ def trapz(y, x, axis):
 )
 
 
+def astype(data, dtype, **kwargs):
+    return data.astype(dtype, **kwargs)
+
+
 def asarray(data, xp=np):
     return (
         data

xarray/core/ops.py

@@ -42,7 +42,7 @@
 NUMPY_SAME_METHODS = ["item", "searchsorted"]
 # methods which don't modify the data shape, so the result should still be
 # wrapped in an Variable/DataArray
-NUMPY_UNARY_METHODS = ["astype", "argsort", "clip", "conj", "conjugate"]
+NUMPY_UNARY_METHODS = ["argsort", "clip", "conj", "conjugate"]
 PANDAS_UNARY_FUNCTIONS = ["isnull", "notnull"]
 # methods which remove an axis
 REDUCE_METHODS = ["all", "any"]

xarray/core/variable.py

@@ -360,6 +360,37 @@ def data(self, data):
             )
         self._data = data
 
+    def astype(self, dtype, casting="unsafe", copy=True):
+        """
+        Copy of the Variable object, with data cast to a specified type.
+
+        Parameters
+        ----------
+        dtype : str or dtype
+             Typecode or data-type to which the array is cast.
+        casting : {'no', 'equiv', 'safe', 'same_kind', 'unsafe'}, optional
+             Controls what kind of data casting may occur. Defaults to 'unsafe'
+             for backwards compatibility.
+
+             * 'no' means the data types should not be cast at all.
+             * 'equiv' means only byte-order changes are allowed.
+             * 'safe' means only casts which can preserve values are allowed.
+             * 'same_kind' means only safe casts or casts within a kind,
+                 like float64 to float32, are allowed.
+             * 'unsafe' means any data conversions may be done.
+        copy : bool, optional
+             By default, astype always returns a newly allocated array. If this
+             is set to False and the `dtype` requirement is satisfied, the input
+             array is returned instead of a copy.
+
+        See also
+        --------
+        np.ndarray.astype
+        dask.array.Array.astype
+        """
+        self.data = duck_array_ops.astype(self.data, dtype, casting=casting, copy=copy)
+        return self
+
     def load(self, **kwargs):
         """Manually trigger loading of this variable's data from disk or a
         remote source into memory and return this variable.

xarray/tests/test_dataarray.py

@@ -1874,6 +1874,13 @@ def test_array_interface(self):
         bar = Variable(["x", "y"], np.zeros((10, 20)))
         assert_equal(self.dv, np.maximum(self.dv, bar))
 
+    def test_astype_attrs(self):
+        mda1 = self.mda.copy()
+        mda1.attrs["foo"] = "bar"
+        mda2 = mda1.astype(bool)
+
+        assert list(mda1.attrs.items()) == list(mda2.attrs.items())
+
     def test_is_null(self):
         x = np.random.RandomState(42).randn(5, 6)
         x[x < 0] = np.nan

xarray/tests/test_dataset.py

@@ -5607,6 +5607,14 @@ def test_pad(self):
         np.testing.assert_equal(padded["var1"].isel(dim2=[0, -1]).data, 42)
         np.testing.assert_equal(padded["dim2"][[0, -1]].data, np.nan)
 
+    def test_astype_attrs(self):
+        data = create_test_data(seed=123)
+        data.attrs["foo"] = "bar"
+        databool = data.astype(bool)
+
+        assert list(data.attrs.items()) == list(databool.attrs.items())
+        assert list(data.var1.attrs.items()) == list(databool.var1.attrs.items())
+
 
 # Py.test tests