Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

keep_attrs for pad #7267

Merged
merged 15 commits into from
Dec 12, 2022
Merged
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