TST: Add Styler.where test

Author: Licht-T

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