ENH: Add Styler.where (pandas-dev#17474)

Author: Licht-T

doc/source/api.rst

@@ -2062,6 +2062,7 @@ Style Application
 
    Styler.apply
    Styler.applymap
+   Styler.where
    Styler.format
    Styler.set_precision
    Styler.set_table_styles

doc/source/whatsnew/v0.21.0.txt

@@ -112,6 +112,7 @@ Other Enhancements
 - `read_*` methods can now infer compression from non-string paths, such as ``pathlib.Path`` objects (:issue:`17206`).
 - :func:`pd.read_sas()` now recognizes much more of the most frequently used date (datetime) formats in SAS7BDAT files (:issue:`15871`).
 - :func:`DataFrame.items` and :func:`Series.items` is now present in both Python 2 and 3 and is lazy in all cases (:issue:`13918`, :issue:`17213`)
+- :func:`Styler.where` has been implemented. It is as a convenience for :func:`Styler.applymap` and enables simple DataFrame styling on the Jupyter notebook (:issue:`17474`).
 
 
 

pandas/io/formats/style.py

@@ -618,11 +618,53 @@ def applymap(self, func, subset=None, **kwargs):
         -------
         self : Styler
 
+        See Also
+        --------
+        Styler.where
+
         """
         self._todo.append((lambda instance: getattr(instance, '_applymap'),
                            (func, subset), kwargs))
         return self
 
+    def where(self, cond, value, other=None, subset=None, **kwargs):
+        """
+        Apply a function elementwise, updating the HTML
+        representation with a style which is selected in
+        accordance with the return value of a function.
+
+        .. versionadded:: 0.21.0
+
+        Parameters
+        ----------
+        cond : callable
+            ``cond`` should take a scalar and return a boolean
+        value : str
+            applied when ``cond`` returns true
+        other : str
+            applied when ``cond`` returns false
+        subset : IndexSlice
+            a valid indexer to limit ``data`` to *before* applying the
+            function. Consider using a pandas.IndexSlice
+        kwargs : dict
+            pass along to ``cond``
+
+        Returns
+        -------
+        self : Styler
+
+        See Also
+        --------
+        Styler.applymap
+
+        """
+
+        if other is None:
+            other = ''
+
+        return self.applymap(lambda val: value if cond(val) else other,
+                             subset=subset, **kwargs)
+
     def set_precision(self, precision):
         """
         Set the precision used to render.

pandas/tests/io/formats/test_style.py

@@ -265,6 +265,64 @@ def f(x):
                             col in self.df.loc[slice_].columns)
             assert result == expected
 
+    def test_where_with_one_style(self):
+        # GH 17474
+        def f(x):
+            return x > 0.5
+
+        style1 = 'foo: bar'
+
+        result = self.df.style.where(f, style1)._compute().ctx
+        expected = dict(((r, c),
+                        [style1 if f(self.df.loc[row, col]) else ''])
+                        for r, row in enumerate(self.df.index)
+                        for c, col in enumerate(self.df.columns))
+        assert result == expected
+
+    def test_where_subset(self):
+        # GH 17474
+        def f(x):
+            return x > 0.5
+
+        style1 = 'foo: bar'
+        style2 = 'baz: foo'
+
+        slices = [pd.IndexSlice[:], pd.IndexSlice[:, ['A']],
+                  pd.IndexSlice[[1], :], pd.IndexSlice[[1], ['A']],
+                  pd.IndexSlice[:2, ['A', 'B']]]
+
+        for slice_ in slices:
+            result = self.df.style.where(f, style1, style2,
+                                         subset=slice_)._compute().ctx
+            expected = dict(((r, c),
+                            [style1 if f(self.df.loc[row, col]) else style2])
+                            for r, row in enumerate(self.df.index)
+                            for c, col in enumerate(self.df.columns)
+                            if row in self.df.loc[slice_].index and
+                            col in self.df.loc[slice_].columns)
+            assert result == expected
+
+    def test_where_subset_compare_with_applymap(self):
+        # GH 17474
+        def f(x):
+            return x > 0.5
+
+        style1 = 'foo: bar'
+        style2 = 'baz: foo'
+
+        def g(x):
+            return style1 if f(x) else style2
+
+        slices = [pd.IndexSlice[:], pd.IndexSlice[:, ['A']],
+                  pd.IndexSlice[[1], :], pd.IndexSlice[[1], ['A']],
+                  pd.IndexSlice[:2, ['A', 'B']]]
+
+        for slice_ in slices:
+            result = self.df.style.where(f, style1, style2,
+                                         subset=slice_)._compute().ctx
+            expected = self.df.style.applymap(g, subset=slice_)._compute().ctx
+            assert result == expected
+
     def test_empty(self):
         df = pd.DataFrame({'A': [1, 0]})
         s = df.style