BUG: Accept generic array-like in .where

[ci skip]

Author: gfyoung

doc/source/whatsnew/v0.20.0.txt

@@ -509,6 +509,7 @@ Bug Fixes
 
 
 
+- Bug in ``Series.where()`` and ``DataFrame.where()`` where array-like conditionals were being rejected (:issue:``)
 - Bug in ``Series`` construction with a datetimetz (:issue:`14928`)
 
 - Bug in compat for passing long integers to ``Timestamp.replace`` (:issue:`15030`)

pandas/core/generic.py

@@ -4719,8 +4719,7 @@ def _where(self, cond, other=np.nan, inplace=False, axis=None, level=None,
             cond, _ = cond.align(self, join='right', broadcast_axis=1)
         else:
             if not hasattr(cond, 'shape'):
-                raise ValueError('where requires an ndarray like object for '
-                                 'its condition')
+                cond = np.asanyarray(cond)
             if cond.shape != self.shape:
                 raise ValueError('Array conditional must be same shape as '
                                  'self')
@@ -4878,26 +4877,20 @@ def _where(self, cond, other=np.nan, inplace=False, axis=None, level=None,
 
         Parameters
         ----------
-        cond : boolean %(klass)s, array or callable
+        cond : boolean %(klass)s, array-like, or callable
             If cond is callable, it is computed on the %(klass)s and
-            should return boolean %(klass)s or array.
-            The callable must not change input %(klass)s
-            (though pandas doesn't check it).
+            should return boolean %(klass)s or array. The callable must
+            not change input %(klass)s (though pandas doesn't check it).
 
             .. versionadded:: 0.18.1
 
-            A callable can be used as cond.
-
         other : scalar, %(klass)s, or callable
             If other is callable, it is computed on the %(klass)s and
-            should return scalar or %(klass)s.
-            The callable must not change input %(klass)s
-            (though pandas doesn't check it).
+            should return scalar or %(klass)s. The callable must not
+            change input %(klass)s (though pandas doesn't check it).
 
             .. versionadded:: 0.18.1
 
-            A callable can be used as other.
-
         inplace : boolean, default False
             Whether to perform the operation in place on the data
         axis : alignment axis if needed, default None

pandas/indexes/base.py

@@ -573,7 +573,7 @@ def repeat(self, repeats, *args, **kwargs):
 
         Parameters
         ----------
-        cond : boolean same length as self
+        cond : boolean array with the same length as self
         other : scalar, or array-like
         """
 

pandas/tests/frame/test_indexing.py

@@ -2479,6 +2479,25 @@ def _check_set(df, cond, check_dtypes=True):
         expected = df[df['a'] == 1].reindex(df.index)
         assert_frame_equal(result, expected)
 
+    def test_where_array_like(self):
+        klasses = [list, tuple, np.array]
+
+        df = DataFrame({'a': [1, 2, 3]})
+        cond = [[False], [True], [True]]
+        expected = DataFrame({'a': [np.nan, 2, 3]})
+
+        for klass in klasses:
+            result = df.where(klass(cond))
+            assert_frame_equal(result, expected)
+
+        df['b'] = 2
+        expected['b'] = [2, np.nan, 2]
+        cond = [[False, True], [True, False], [True, True]]
+
+        for klass in klasses:
+            result = df.where(klass(cond))
+            assert_frame_equal(result, expected)
+
     def test_where_bug(self):
 
         # GH 2793

pandas/tests/indexes/common.py

@@ -497,6 +497,18 @@ def test_where(self):
         result = i.where(cond)
         tm.assert_index_equal(result, expected)
 
+    def test_where_array_like(self):
+        i = self.create_index()
+
+        _nan = i._na_value
+        cond = [False] + [True] * (len(i) - 1)
+        klasses = [list, tuple, np.array, pd.Series]
+        expected = pd.Index([_nan] + i[1:].tolist(), dtype=i.dtype)
+
+        for klass in klasses:
+            result = i.where(klass(cond))
+            tm.assert_index_equal(result, expected)
+
     def test_setops_errorcases(self):
         for name, idx in compat.iteritems(self.indices):
             # # non-iterable input

pandas/tests/indexes/period/test_period.py

@@ -89,13 +89,22 @@ def test_where(self):
         expected = i
         tm.assert_index_equal(result, expected)
 
-        i2 = i.copy()
         i2 = pd.PeriodIndex([pd.NaT, pd.NaT] + i[2:].tolist(),
                             freq='D')
         result = i.where(notnull(i2))
         expected = i2
         tm.assert_index_equal(result, expected)
 
+    def test_where_array_like(self):
+        i = self.create_index()
+        cond = [False] + [True] * (len(i) - 1)
+        klasses = [list, tuple, np.array, Series]
+        expected = pd.PeriodIndex([pd.NaT] + i[1:].tolist(), freq='D')
+
+        for klass in klasses:
+            result = i.where(klass(cond))
+            tm.assert_index_equal(result, expected)
+
     def test_where_other(self):
 
         i = self.create_index()

pandas/tests/indexes/test_category.py

@@ -240,13 +240,23 @@ def test_where(self):
         expected = i
         tm.assert_index_equal(result, expected)
 
-        i2 = i.copy()
         i2 = pd.CategoricalIndex([np.nan, np.nan] + i[2:].tolist(),
                                  categories=i.categories)
         result = i.where(notnull(i2))
         expected = i2
         tm.assert_index_equal(result, expected)
 
+    def test_where_array_like(self):
+        i = self.create_index()
+        cond = [False] + [True] * (len(i) - 1)
+        klasses = [list, tuple, np.array, pd.Series]
+        expected = pd.CategoricalIndex([np.nan] + i[1:].tolist(),
+                                       categories=i.categories)
+
+        for klass in klasses:
+            result = i.where(klass(cond))
+            tm.assert_index_equal(result, expected)
+
     def test_append(self):
 
         ci = self.create_index()

pandas/tests/indexes/test_multi.py

@@ -88,6 +88,15 @@ def f():
 
         self.assertRaises(NotImplementedError, f)
 
+    def test_where_array_like(self):
+        i = MultiIndex.from_tuples([('A', 1), ('A', 2)])
+        klasses = [list, tuple, np.array, pd.Series]
+        cond = [False, True]
+
+        for klass in klasses:
+            f = lambda: i.where(klass(cond))
+            self.assertRaises(NotImplementedError, f)
+
     def test_repeat(self):
         reps = 2
         numbers = [1, 2, 3]

pandas/tests/series/test_indexing.py

@@ -1193,6 +1193,16 @@ def f():
         expected = Series(np.nan, index=[9])
         assert_series_equal(result, expected)
 
+    def test_where_array_like(self):
+        s = Series([1, 2, 3])
+        cond = [False, True, True]
+        expected = Series([np.nan, 2, 3])
+        klasses = [list, tuple, np.array, Series]
+
+        for klass in klasses:
+            result = s.where(klass(cond))
+            assert_series_equal(result, expected)
+
     def test_where_setitem_invalid(self):
 
         # GH 2702