add where and mask methods to Series. where returns a series evaluated for the cond with a shape like the original

Author: jreback

pandas/core/series.py

@@ -562,6 +562,44 @@ def _get_values(self, indexer):
         except Exception:
             return self.values[indexer]
 
+    def where(self, cond, other=nan, inplace=False):
+        """
+        Return a Series where cond is True; otherwise values are from other
+
+        Parameters
+        ----------
+        cond: boolean Series or array
+        other: scalar or Series
+
+        Returns
+        -------
+        wh: Series
+        """
+        if not hasattr(cond, 'shape'):
+            raise ValueError('where requires an ndarray like object for its '
+                             'condition')
+
+        if inplace:
+            self._set_with(~cond, other)
+            return self
+
+        return self._get_values(cond).reindex_like(self).fillna(other)
+
+    def mask(self, cond):
+        """
+        Returns copy of self whose values are replaced with nan if the
+        inverted condition is True
+
+        Parameters
+        ----------
+        cond: boolean Series or array
+
+        Returns
+        -------
+        wh: Series
+        """
+        return self.where(~cond, nan)
+
     def __setitem__(self, key, value):
         try:
             try:

pandas/tests/test_series.py

@@ -939,6 +939,37 @@ def test_ix_getitem_iterator(self):
         result = self.series.ix[idx]
         assert_series_equal(result, self.series[:10])
 
+    def test_where(self):
+        s = Series(np.random.randn(5))
+        cond = s > 0
+
+        rs  = s.where(cond).dropna()
+        rs2 = s[cond]
+        assert_series_equal(rs, rs2)
+
+        rs  = s.where(cond,-s)
+        assert_series_equal(rs, s.abs())
+
+        rs  = s.where(cond)
+        assert(s.shape == rs.shape)
+
+        self.assertRaises(ValueError, s.where, 1)
+
+    def test_where_inplace(self):
+        s = Series(np.random.randn(5))
+        cond = s > 0
+
+        rs = s.copy()
+        rs.where(cond,inplace=True)
+        assert_series_equal(rs.dropna(), s[cond])
+
+    def test_mask(self):
+        s = Series(np.random.randn(5))
+        cond = s > 0
+
+        rs = s.where(cond, np.nan)
+        assert_series_equal(rs, s.mask(~cond))
+
     def test_ix_setitem(self):
         inds = self.series.index[[3,4,7]]