diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 2372b873422..e2f0cfcc5f6 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -36,6 +36,9 @@ Bug fixes ~~~~~~~~~ - Allow numpy-only objects in :py:func:`where` when ``keep_attrs=True`` (:issue:`7362`, :pull:`7364`). By `Sam Levang `_. +- add a ``keep_attrs`` parameter to :py:meth:`Dataset.pad`, :py:meth:`DataArray.pad`, + and :py:meth:`Variable.pad` (:pull:`7267`). + By `Justus Magin `_. Documentation ~~~~~~~~~~~~~ diff --git a/xarray/core/dataarray.py b/xarray/core/dataarray.py index 273d1027283..dc8dcbd3776 100644 --- a/xarray/core/dataarray.py +++ b/xarray/core/dataarray.py @@ -5270,6 +5270,7 @@ def pad( | None = None, end_values: int | tuple[int, int] | Mapping[Any, tuple[int, int]] | None = None, reflect_type: PadReflectOptions = None, + keep_attrs: bool | None = None, **pad_width_kwargs: Any, ) -> T_DataArray: """Pad this array along one or more dimensions. @@ -5347,6 +5348,10 @@ def pad( default with an unaltered reflection around the edge value. For the "odd" style, the extended part of the array is created by subtracting the reflected values from two times the edge value. + keep_attrs : bool or None, optional + If True, the attributes (``attrs``) will be copied from the + original object to the new one. If False, the new object + will be returned without attributes. **pad_width_kwargs The keyword arguments form of ``pad_width``. One of ``pad_width`` or ``pad_width_kwargs`` must be provided. @@ -5414,6 +5419,7 @@ def pad( constant_values=constant_values, end_values=end_values, reflect_type=reflect_type, + keep_attrs=keep_attrs, **pad_width_kwargs, ) return self._from_temp_dataset(ds) diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index 2e2fd6efa72..3d4bf0a9a20 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -7939,6 +7939,7 @@ def pad( ) = None, end_values: int | tuple[int, int] | Mapping[Any, tuple[int, int]] | None = None, reflect_type: PadReflectOptions = None, + keep_attrs: bool | None = None, **pad_width_kwargs: Any, ) -> T_Dataset: """Pad this dataset along one or more dimensions. @@ -8016,6 +8017,10 @@ def pad( default with an unaltered reflection around the edge value. For the "odd" style, the extended part of the array is created by subtracting the reflected values from two times the edge value. + keep_attrs : bool or None, optional + If True, the attributes (``attrs``) will be copied from the + original object to the new one. If False, the new object + will be returned without attributes. **pad_width_kwargs The keyword arguments form of ``pad_width``. One of ``pad_width`` or ``pad_width_kwargs`` must be provided. @@ -8062,6 +8067,9 @@ def pad( coord_pad_mode = "constant" coord_pad_options = {} + if keep_attrs is None: + keep_attrs = _get_keep_attrs(default=True) + variables = {} # keep indexes that won't be affected by pad and drop all other indexes @@ -8084,11 +8092,13 @@ def pad( constant_values=constant_values, end_values=end_values, reflect_type=reflect_type, + keep_attrs=keep_attrs, ) else: variables[name] = var.pad( pad_width=var_pad_width, mode=coord_pad_mode, + keep_attrs=keep_attrs, **coord_pad_options, # type: ignore[arg-type] ) # reset default index of dimension coordinates @@ -8099,7 +8109,8 @@ def pad( indexes[name] = index variables[name] = index_vars[name] - return self._replace_with_new_dims(variables, indexes=indexes) + attrs = self._attrs if keep_attrs else None + return self._replace_with_new_dims(variables, indexes=indexes, attrs=attrs) def idxmin( self: T_Dataset, diff --git a/xarray/core/variable.py b/xarray/core/variable.py index d3db3fb2dd6..27d716b9f93 100644 --- a/xarray/core/variable.py +++ b/xarray/core/variable.py @@ -1432,6 +1432,7 @@ def pad( | None = None, end_values: int | tuple[int, int] | Mapping[Any, tuple[int, int]] | None = None, reflect_type: PadReflectOptions = None, + keep_attrs: bool | None = None, **pad_width_kwargs: Any, ): """ @@ -1459,6 +1460,10 @@ def pad( default with an unaltered reflection around the edge value. For the "odd" style, the extended part of the array is created by subtracting the reflected values from two times the edge value. + keep_attrs : bool, optional + If True, the variable's attributes (`attrs`) will be copied from + the original object to the new one. If False (default), the new + object will be returned without attributes. **pad_width_kwargs One of pad_width or pad_width_kwargs must be provided. @@ -1515,7 +1520,11 @@ def pad( **pad_option_kwargs, ) - return type(self)(self.dims, array) + if keep_attrs is None: + keep_attrs = _get_keep_attrs(default=True) + attrs = self._attrs if keep_attrs else None + + return type(self)(self.dims, array, attrs=attrs) def _roll_one_dim(self, dim, count): axis = self.get_axis_num(dim) diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py index 8184fe1955c..8a1cfea55a0 100644 --- a/xarray/tests/test_dataarray.py +++ b/xarray/tests/test_dataarray.py @@ -4166,6 +4166,36 @@ def test_pad_reflect(self, mode, reflect_type) -> None: assert actual.shape == (7, 4, 9) assert_identical(actual, expected) + @pytest.mark.parametrize( + ["keep_attrs", "attrs", "expected"], + [ + pytest.param(None, {"a": 1, "b": 2}, {"a": 1, "b": 2}, id="default"), + pytest.param(False, {"a": 1, "b": 2}, {}, id="False"), + pytest.param(True, {"a": 1, "b": 2}, {"a": 1, "b": 2}, id="True"), + ], + ) + def test_pad_keep_attrs(self, keep_attrs, attrs, expected) -> None: + arr = xr.DataArray( + [1, 2], dims="x", coords={"c": ("x", [-1, 1], attrs)}, attrs=attrs + ) + expected = xr.DataArray( + [0, 1, 2, 0], + dims="x", + coords={"c": ("x", [np.nan, -1, 1, np.nan], expected)}, + attrs=expected, + ) + + keep_attrs_ = "default" if keep_attrs is None else keep_attrs + + with set_options(keep_attrs=keep_attrs_): + actual = arr.pad({"x": (1, 1)}, mode="constant", constant_values=0) + xr.testing.assert_identical(actual, expected) + + actual = arr.pad( + {"x": (1, 1)}, mode="constant", constant_values=0, keep_attrs=keep_attrs + ) + xr.testing.assert_identical(actual, expected) + @pytest.mark.parametrize("parser", ["pandas", "python"]) @pytest.mark.parametrize( "engine", ["python", None, pytest.param("numexpr", marks=[requires_numexpr])] diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py index 0d3be9d378b..6ab4c40cfa7 100644 --- a/xarray/tests/test_dataset.py +++ b/xarray/tests/test_dataset.py @@ -6115,6 +6115,40 @@ def test_pad(self) -> None: np.testing.assert_equal(padded["var1"].isel(dim2=[0, -1]).data, 42) np.testing.assert_equal(padded["dim2"][[0, -1]].data, np.nan) + @pytest.mark.parametrize( + ["keep_attrs", "attrs", "expected"], + [ + pytest.param(None, {"a": 1, "b": 2}, {"a": 1, "b": 2}, id="default"), + pytest.param(False, {"a": 1, "b": 2}, {}, id="False"), + pytest.param(True, {"a": 1, "b": 2}, {"a": 1, "b": 2}, id="True"), + ], + ) + def test_pad_keep_attrs(self, keep_attrs, attrs, expected) -> None: + ds = xr.Dataset( + {"a": ("x", [1, 2], attrs), "b": ("y", [1, 2], attrs)}, + coords={"c": ("x", [-1, 1], attrs), "d": ("y", [-1, 1], attrs)}, + attrs=attrs, + ) + expected = xr.Dataset( + {"a": ("x", [0, 1, 2, 0], expected), "b": ("y", [1, 2], attrs)}, + coords={ + "c": ("x", [np.nan, -1, 1, np.nan], expected), + "d": ("y", [-1, 1], attrs), + }, + attrs=expected, + ) + + keep_attrs_ = "default" if keep_attrs is None else keep_attrs + + with set_options(keep_attrs=keep_attrs_): + actual = ds.pad({"x": (1, 1)}, mode="constant", constant_values=0) + xr.testing.assert_identical(actual, expected) + + actual = ds.pad( + {"x": (1, 1)}, mode="constant", constant_values=0, keep_attrs=keep_attrs + ) + xr.testing.assert_identical(actual, expected) + def test_astype_attrs(self) -> None: data = create_test_data(seed=123) data.attrs["foo"] = "bar" diff --git a/xarray/tests/test_variable.py b/xarray/tests/test_variable.py index c82a8c3a539..8eadf7d51fd 100644 --- a/xarray/tests/test_variable.py +++ b/xarray/tests/test_variable.py @@ -910,6 +910,33 @@ def test_pad_constant_values(self, xr_arg, np_arg): ) assert_array_equal(actual, expected) + @pytest.mark.parametrize( + ["keep_attrs", "attrs", "expected"], + [ + pytest.param(None, {"a": 1, "b": 2}, {"a": 1, "b": 2}, id="default"), + pytest.param(False, {"a": 1, "b": 2}, {}, id="False"), + pytest.param(True, {"a": 1, "b": 2}, {"a": 1, "b": 2}, id="True"), + ], + ) + def test_pad_keep_attrs(self, keep_attrs, attrs, expected): + data = np.arange(10, dtype=float) + v = self.cls(["x"], data, attrs) + + keep_attrs_ = "default" if keep_attrs is None else keep_attrs + + with set_options(keep_attrs=keep_attrs_): + actual = v.pad({"x": (1, 1)}, mode="constant", constant_values=np.nan) + + assert actual.attrs == expected + + actual = v.pad( + {"x": (1, 1)}, + mode="constant", + constant_values=np.nan, + keep_attrs=keep_attrs, + ) + assert actual.attrs == expected + @pytest.mark.parametrize("d, w", (("x", 3), ("y", 5))) def test_rolling_window(self, d, w): # Just a working test. See test_nputils for the algorithm validation