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

allow other and drop arguments in where (gh#6466) #6467

Merged
merged 6 commits into from
Apr 12, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions doc/whats-new.rst
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,9 @@ Bug fixes
By `Stan West <https://github.com/stanwest>`_.
- Fix bug in :py:func:`where` when passing non-xarray objects with ``keep_attrs=True``. (:issue:`6444`, :pull:`6461`)
By `Sam Levang <https://github.com/slevang>`_.
- Allow passing both ``other`` and ``drop=True`` arguments to ``xr.DataArray.where``
and ``xr.Dataset.where`` (:pull:`6466`, :pull:`6467`).
By `Michael Delgado <https://github.com/delgadom>`_.

Documentation
~~~~~~~~~~~~~
Expand Down
14 changes: 9 additions & 5 deletions xarray/core/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -1197,8 +1197,7 @@ def where(self, cond, other=dtypes.NA, drop: bool = False):
By default, these locations filled with NA.
drop : bool, optional
If True, coordinate labels that only correspond to False values of
the condition are dropped from the result. Mutually exclusive with
``other``.
the condition are dropped from the result.

Returns
-------
Expand Down Expand Up @@ -1251,6 +1250,14 @@ def where(self, cond, other=dtypes.NA, drop: bool = False):
[15., nan, nan, nan]])
Dimensions without coordinates: x, y

>>> a.where(a.x + a.y < 4, -1, drop=True)
<xarray.DataArray (x: 4, y: 4)>
array([[ 0, 1, 2, 3],
[ 5, 6, 7, -1],
[10, 11, -1, -1],
[15, -1, -1, -1]])
Dimensions without coordinates: x, y

See Also
--------
numpy.where : corresponding numpy function
Expand All @@ -1264,9 +1271,6 @@ def where(self, cond, other=dtypes.NA, drop: bool = False):
cond = cond(self)

if drop:
if other is not dtypes.NA:
raise ValueError("cannot set `other` if drop=True")

if not isinstance(cond, (Dataset, DataArray)):
raise TypeError(
f"cond argument is {cond!r} but must be a {Dataset!r} or {DataArray!r}"
Expand Down
7 changes: 5 additions & 2 deletions xarray/tests/test_dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -4548,8 +4548,11 @@ def test_where_other(self):
actual = ds.where(lambda x: x > 1, -1)
assert_equal(expected, actual)

with pytest.raises(ValueError, match=r"cannot set"):
ds.where(ds > 1, other=0, drop=True)
actual = ds.where(ds > 1, other=-1, drop=True)
expected_nodrop = ds.where(ds > 1, -1)
_, expected = xr.align(actual, expected_nodrop, join="left")
assert_equal(actual, expected)
assert actual.a.dtype == int

with pytest.raises(ValueError, match=r"cannot align .* are not equal"):
ds.where(ds > 1, ds.isel(x=slice(3)))
Expand Down