From 72874df5d87e60915b996af8580f910a0c49b7e2 Mon Sep 17 00:00:00 2001 From: Michael Delgado Date: Sun, 10 Apr 2022 10:58:24 -0700 Subject: [PATCH 1/4] allow other and drop arguments in where (gh#6466) --- xarray/core/common.py | 14 +++++++++----- xarray/tests/test_dataset.py | 7 +++++-- 2 files changed, 14 insertions(+), 7 deletions(-) diff --git a/xarray/core/common.py b/xarray/core/common.py index c33db4a62ea..817b63161bc 100644 --- a/xarray/core/common.py +++ b/xarray/core/common.py @@ -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 ------- @@ -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) + + 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 @@ -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}" diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py index 5f368375fc0..45aa82b6fc7 100644 --- a/xarray/tests/test_dataset.py +++ b/xarray/tests/test_dataset.py @@ -4545,8 +4545,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))) From 239b0f01f1f0f5fb807db9e15a11fbdc735fab14 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 10 Apr 2022 18:05:09 +0000 Subject: [PATCH 2/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- xarray/tests/test_dataset.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py index 45aa82b6fc7..64fb46daa08 100644 --- a/xarray/tests/test_dataset.py +++ b/xarray/tests/test_dataset.py @@ -4547,7 +4547,7 @@ def test_where_other(self): actual = ds.where(ds > 1, other=-1, drop=True) expected_nodrop = ds.where(ds > 1, -1) - _, expected = xr.align(actual, expected_nodrop, join='left') + _, expected = xr.align(actual, expected_nodrop, join="left") assert_equal(actual, expected) assert actual.a.dtype == int From 9fc6c1b4a7d68d0a1c6e06a753173c11fd720cd4 Mon Sep 17 00:00:00 2001 From: Michael Delgado Date: Sun, 10 Apr 2022 11:06:59 -0700 Subject: [PATCH 3/4] add whatsnew entry to bug fixes? --- doc/whats-new.rst | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index c84a0549774..77a67068edf 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -71,8 +71,11 @@ Bug fixes - Allow fancy indexing of duck dask arrays along multiple dimensions. (:pull:`6414`) By `Justus Magin `_. - In the API for backends, support dimensions that express their preferred chunk sizes - as a tuple of integers. (:issue:`6333`, :pull:`6334`) + as a tuple of integers. (:issue:`6333`, :pull:`6334`). By `Stan West `_. +- Allow passing both ``other`` and ``drop=True`` arguments to ``xr.DataArray.where`` + and ``xr.Dataset.where`` (:pull:`6466`, :pull:`6467`). + By `Michael Delgado `_. Documentation ~~~~~~~~~~~~~ From 763b6528fa4b62724ca5f44aae307cfdfa91c0ec Mon Sep 17 00:00:00 2001 From: Michael Delgado Date: Sun, 10 Apr 2022 11:11:48 -0700 Subject: [PATCH 4/4] drop extraneous edit --- doc/whats-new.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 77a67068edf..db7266806c6 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -71,7 +71,7 @@ Bug fixes - Allow fancy indexing of duck dask arrays along multiple dimensions. (:pull:`6414`) By `Justus Magin `_. - In the API for backends, support dimensions that express their preferred chunk sizes - as a tuple of integers. (:issue:`6333`, :pull:`6334`). + as a tuple of integers. (:issue:`6333`, :pull:`6334`) By `Stan West `_. - Allow passing both ``other`` and ``drop=True`` arguments to ``xr.DataArray.where`` and ``xr.Dataset.where`` (:pull:`6466`, :pull:`6467`).