allow where to receive a callable

Author: max-sixty

xarray/core/common.py

@@ -1152,6 +1152,9 @@ def where(self, cond, other=dtypes.NA, drop: bool = False):
         from .dataarray import DataArray
         from .dataset import Dataset
 
+        if callable(cond):
+            return self.where(cond(self), other=other, drop=drop)
+
         if drop:
             if other is not dtypes.NA:
                 raise ValueError("cannot set `other` if drop=True")

xarray/tests/test_dataarray.py

@@ -2215,6 +2215,12 @@ def test_where(self):
         actual = arr.where(arr.x < 2, drop=True)
         assert_identical(actual, expected)
 
+    def test_where_lambda(self):
+        arr = DataArray(np.arange(4), dims="y")
+        expected = arr.sel(y=slice(2))
+        actual = arr.where(lambda x: x.y < 2, drop=True)
+        assert_identical(actual, expected)
+
     def test_where_string(self):
         array = DataArray(["a", "b"])
         expected = DataArray(np.array(["a", np.nan], dtype=object))

xarray/tests/test_dataset.py

@@ -4349,13 +4349,22 @@ def test_where(self):
         assert actual.a.name == "a"
         assert actual.a.attrs == ds.a.attrs
 
+        # lambda
+        ds = Dataset({"a": ("x", range(5))})
+        expected = Dataset({"a": ("x", [np.nan, np.nan, 2, 3, 4])})
+        actual = ds.where(lambda x: x > 1)
+        assert_identical(expected, actual)
+
     def test_where_other(self):
         ds = Dataset({"a": ("x", range(5))}, {"x": range(5)})
         expected = Dataset({"a": ("x", [-1, -1, 2, 3, 4])}, {"x": range(5)})
         actual = ds.where(ds > 1, -1)
         assert_equal(expected, actual)
         assert actual.a.dtype == int
 
+        actual = ds.where(lambda x: x > 1, -1)
+        assert_equal(expected, actual)
+
         with raises_regex(ValueError, "cannot set"):
             ds.where(ds > 1, other=0, drop=True)