Skip to content

Commit

Permalink
keep_attrs for pad (#7267)
Browse files Browse the repository at this point in the history
* make use of `keep_attrs` in `Variable.pad`

* propagate `keep_attrs`

* resolve the default `keep_attrs` in `Dataset.pad`

* add tests for `Variable.pad`

* add tests for `Dataset.pad` and keep the `Dataset` attrs

* actually pad the variable

* also check that untouched variables are not modified

* add tests for `DataArray.pad`

* update whats-new.rst

* Set True by default

* move the whats-new.rst entry to the dev section

Co-authored-by: Deepak Cherian <dcherian@users.noreply.github.com>
Co-authored-by: Illviljan <14371165+Illviljan@users.noreply.github.com>
  • Loading branch information
3 people authored Dec 12, 2022
1 parent b610956 commit db68db6
Show file tree
Hide file tree
Showing 7 changed files with 122 additions and 2 deletions.
3 changes: 3 additions & 0 deletions doc/whats-new.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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 <https://github.com/slevang>`_.
- add a ``keep_attrs`` parameter to :py:meth:`Dataset.pad`, :py:meth:`DataArray.pad`,
and :py:meth:`Variable.pad` (:pull:`7267`).
By `Justus Magin <https://github.com/keewis>`_.

Documentation
~~~~~~~~~~~~~
Expand Down
6 changes: 6 additions & 0 deletions xarray/core/dataarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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)
Expand Down
13 changes: 12 additions & 1 deletion xarray/core/dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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,
Expand Down
11 changes: 10 additions & 1 deletion xarray/core/variable.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
):
"""
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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)
Expand Down
30 changes: 30 additions & 0 deletions xarray/tests/test_dataarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -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])]
Expand Down
34 changes: 34 additions & 0 deletions xarray/tests/test_dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
27 changes: 27 additions & 0 deletions xarray/tests/test_variable.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit db68db6

Please sign in to comment.