added implementation + tests for scalar condition in .where

Author: randolf-scholz

pandas/core/generic.py

@@ -10102,7 +10102,12 @@ def _where(
         else:
             if not hasattr(cond, "shape"):
                 cond = np.asanyarray(cond)
-            if cond.shape != () and cond.shape != self.shape:
+            if cond.shape == ():
+                # Note: DataFrame(True, index=[1,2,3], columns=["a", "b", "c"]) works
+                # but DataFrame(np.array(True), index=[1,2,3], columns=["a", "b", "c"]) does not
+                # hence we need to unpack scalar
+                cond = cond.item()
+            elif cond.shape != self.shape:
                 raise ValueError("Array conditional must be same shape as self")
             cond = self._constructor(cond, **self._construct_axes_dict(), copy=False)
 

pandas/tests/frame/indexing/test_where.py

@@ -159,10 +159,15 @@ def test_where_invalid(self):
         with pytest.raises(ValueError, match=msg):
             df.where(err2, other1)
 
-        with pytest.raises(ValueError, match=msg):
-            df.mask(True)
-        with pytest.raises(ValueError, match=msg):
-            df.mask(0)
+    def test_where_scalar_cond(self):
+        df = DataFrame(np.random.randn(5, 3), columns=["A", "B", "C"])
+        result = df.where(True)
+        expected = df
+        tm.assert_frame_equal(result, expected)
+
+        result = df.where(False)
+        expected = DataFrame(np.nan, index=df.index, columns=df.columns)
+        tm.assert_frame_equal(result, expected)
 
     def test_where_set(self, where_frame, float_string_frame):
         # where inplace

pandas/tests/series/indexing/test_where.py

@@ -147,8 +147,6 @@ def test_where_error():
     cond = s > 0
 
     msg = "Array conditional must be same shape as self"
-    with pytest.raises(ValueError, match=msg):
-        s.where(1)
     with pytest.raises(ValueError, match=msg):
         s.where(cond[:3].values, -s)
 
@@ -466,7 +464,7 @@ def test_where_datetimelike_categorical(tz_naive_fixture):
     tm.assert_frame_equal(res, pd.DataFrame(dr))
 
 
-def test_where_scalar_cond(self):
+def test_where_scalar_cond():
     # True
     ser = Series(pd.Categorical(["a", "b"]))
     result = ser.where(True)
@@ -477,4 +475,4 @@ def test_where_scalar_cond(self):
     ser = Series(pd.Categorical(["a", "b"]))
     result = ser.where(False)
     expected = Series(pd.Categorical([None, None], categories=["a", "b"]))
-    tm.assert_series_equal(result, expected)
+    tm.assert_series_equal(result, expected)