ENH: Styler.bar accepts matplotlib colormap

doc/source/user_guide/style.ipynb

@@ -49,6 +49,7 @@
    "source": [
     "import pandas as pd\n",
     "import numpy as np\n",
+    "import matplotlib as mpl\n",
     "\n",
     "df = pd.DataFrame([[38.0, 2.0, 18.0, 22.0, 21, np.nan],[19, 439, 6, 452, 226,232]], \n",
     "                  index=pd.Index(['Tumour (Positive)', 'Non-Tumour (Negative)'], name='Actual Label:'), \n",
@@ -1275,9 +1276,9 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Additional keyword arguments give more control on centering and positioning, and you can pass a list of `[color_negative, color_positive]` to highlight lower and higher values.\n",
+    "Additional keyword arguments give more control on centering and positioning, and you can pass a list of `[color_negative, color_positive]` to highlight lower and higher values or a matplotlib colormap.\n",
     "\n",
-    "Here's how you can change the above with the new `align` option, combined with setting `vmin` and `vmax` limits, the `width` of the figure, and underlying css `props` of cells, leaving space to display the text and the bars:"
+    "To showcase an example here's how you can change the above with the new `align` option, combined with setting `vmin` and `vmax` limits, the `width` of the figure, and underlying css `props` of cells, leaving space to display the text and the bars. We also use `text_gradient` to color the text the same as the bars using a matplotlib colormap (although in this case the visualization is probably better without this additional effect)."
    ]
   },
   {
@@ -1286,8 +1287,10 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "df2.style.bar(align=0, vmin=-2.5, vmax=2.5, color=['#d65f5f', '#5fba7d'], height=50,\n",
-    "              width=60, props=\"width: 120px; border-right: 1px solid black;\").format('{:.3f}', na_rep=\"\")"
+    "df2.style.format('{:.3f}', na_rep=\"\")\\\n",
+    "         .bar(align=0, vmin=-2.5, vmax=2.5, color=mpl.cm.get_cmap(\"bwr\"), height=50,\n",
+    "              width=60, props=\"width: 120px; border-right: 1px solid black;\")\\\n",
+    "         .text_gradient(cmap=\"bwr\", vmin=-2.5, vmax=2.5)"
    ]
   },
   {
@@ -2031,7 +2034,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.7"
+   "version": "3.8.6"
   }
  },
  "nbformat": 4,

doc/source/whatsnew/v1.4.0.rst

@@ -73,7 +73,7 @@ Styler
 :class:`.Styler` has been further developed in 1.4.0. The following enhancements have been made:
 
   - Styling and formatting of indexes has been added, with :meth:`.Styler.apply_index`, :meth:`.Styler.applymap_index` and :meth:`.Styler.format_index`. These mirror the signature of the methods already used to style and format data values, and work with both HTML and LaTeX format (:issue:`41893`, :issue:`43101`).
-  - :meth:`.Styler.bar` introduces additional arguments to control alignment and display (:issue:`26070`, :issue:`36419`), and it also validates the input arguments ``width`` and ``height`` (:issue:`42511`).
+  - :meth:`.Styler.bar` introduces additional arguments to control alignment, display and colors (:issue:`26070`, :issue:`36419`, :issue:`43662`), and it also validates the input arguments ``width`` and ``height`` (:issue:`42511`).
   - :meth:`.Styler.to_latex` introduces keyword argument ``environment``, which also allows a specific "longtable" entry through a separate jinja2 template (:issue:`41866`).
   - :meth:`.Styler.to_html` introduces keyword arguments ``sparse_index``, ``sparse_columns``, ``bold_headers``, ``caption``, ``max_rows`` and ``max_columns`` (:issue:`41946`, :issue:`43149`, :issue:`42972`).
   - Keyword arguments ``level`` and ``names`` added to :meth:`.Styler.hide_index` and :meth:`.Styler.hide_columns` for additional control of visibility of MultiIndexes and index names (:issue:`25475`, :issue:`43404`, :issue:`43346`)

pandas/io/formats/style.py

@@ -35,7 +35,6 @@
     IndexSlice,
     RangeIndex,
 )
-from pandas.api.types import is_list_like
 from pandas.core import generic
 import pandas.core.common as com
 from pandas.core.frame import (
@@ -2629,11 +2628,12 @@ def bar(
             Apply to each column (``axis=0`` or ``'index'``), to each row
             (``axis=1`` or ``'columns'``), or to the entire DataFrame at once
             with ``axis=None``.
-        color : str or 2-tuple/list
+        color : str, 2-tuple/list, matplotlib Colormap
             If a str is passed, the color is the same for both
             negative and positive numbers. If 2-tuple/list is used, the
             first element is the color_negative and the second is the
-            color_positive (eg: ['#d65f5f', '#5fba7d']).
+            color_positive (eg: ['#d65f5f', '#5fba7d']). Alternatively, assigns
+            colors from a Colormap based on the datavalues.
         width : float, default 100
             The percentage of the cell, measured from the left, in which to draw the
             bars, in [0, 100].
@@ -2676,18 +2676,13 @@ def bar(
         Returns
         -------
         self : Styler
-        """
-        if not (is_list_like(color)):
-            color = [color, color]
-        elif len(color) == 1:
-            color = [color[0], color[0]]
-        elif len(color) > 2:
-            raise ValueError(
-                "`color` must be string or a list-like "
-                "of length 2: [`color_neg`, `color_pos`] "
-                "(eg: color=['#d65f5f', '#5fba7d'])"
-            )
 
+        Notes
+        -----
+        This section of the user guide:
+        `Table Visualization <../../user_guide/style.ipynb>`_ gives
+        a number of examples for different settings and color coordination.
+        """
         if not (0 <= width <= 100):
             raise ValueError(f"`width` must be a value in [0, 100], got {width}")
         elif not (0 <= height <= 100):
@@ -3367,7 +3362,7 @@ def _highlight_value(data: FrameOrSeries, op: str, props: str) -> np.ndarray:
 def _bar(
     data: FrameOrSeries,
     align: str | float | int | Callable,
-    colors: list[str],
+    colors: Any,
     width: float,
     height: float,
     vmin: float | None,
@@ -3429,7 +3424,7 @@ def css_bar(start: float, end: float, color: str) -> str:
             cell_css += f" {color} {end*100:.1f}%, transparent {end*100:.1f}%)"
         return cell_css
 
-    def css_calc(x, left: float, right: float, align: str):
+    def css_calc(x, left: float, right: float, align: str, color: list | str):
         """
         Return the correct CSS for bar placement based on calculated values.
 
@@ -3460,7 +3455,10 @@ def css_calc(x, left: float, right: float, align: str):
         if pd.isna(x):
             return base_css
 
-        color = colors[0] if x < 0 else colors[1]
+        if isinstance(color, list):
+            color = color[0] if x < 0 else color[1]
+        assert isinstance(color, str)  # mypy redefinition
+
         x = left if x < left else x
         x = right if x > right else x  # trim data if outside of the window
 
@@ -3523,15 +3521,48 @@ def css_calc(x, left: float, right: float, align: str):
             "value defining the center line or a callable that returns a float"
         )
 
+    rgbas = None
+    if not isinstance(colors, (list, str)):
+        # use the matplotlib colormap input
+        with _mpl(Styler.bar) as (plt, mpl_colors):
+            norm = mpl_colors.Normalize(left, right)
+            if not isinstance(colors, mpl_colors.Colormap):
+                raise ValueError(
+                    "`colors` must be a matplotlib Colormap if not string "
+                    "or list of strings."
+                )
+            rgbas = colors(norm(values))
+            if data.ndim == 1:
+                rgbas = [mpl_colors.rgb2hex(rgba) for rgba in rgbas]
+            else:
+                rgbas = [[mpl_colors.rgb2hex(rgba) for rgba in row] for row in rgbas]
+    elif isinstance(colors, list) and len(colors) > 2:
+        raise ValueError(
+            "`color` must be string, list-like of 2 strings, or matplotlib Colormap "
+            "(eg: color=['#d65f5f', '#5fba7d'])"
+        )
+
     assert isinstance(align, str)  # mypy: should now be in [left, right, mid, zero]
     if data.ndim == 1:
-        return [css_calc(x - z, left - z, right - z, align) for x in values]
+        return [
+            css_calc(
+                x - z, left - z, right - z, align, colors if rgbas is None else rgbas[i]
+            )
+            for i, x in enumerate(values)
+        ]
     else:
-        return DataFrame(
+        return np.array(
             [
-                [css_calc(x - z, left - z, right - z, align) for x in row]
-                for row in values
-            ],
-            index=data.index,
-            columns=data.columns,
+                [
+                    css_calc(
+                        x - z,
+                        left - z,
+                        right - z,
+                        align,
+                        colors if rgbas is None else rgbas[i][j],
+                    )
+                    for j, x in enumerate(row)
+                ]
+                for i, row in enumerate(values)
+            ]
         )

pandas/tests/io/formats/style/test_matplotlib.py

@@ -10,6 +10,8 @@
 pytest.importorskip("matplotlib")
 pytest.importorskip("jinja2")
 
+import matplotlib as mpl
+
 from pandas.io.formats.style import Styler
 
 
@@ -256,3 +258,26 @@ def test_background_gradient_gmap_wrong_series(styler_blank):
     gmap = Series([1, 2], index=["X", "Y"])
     with pytest.raises(ValueError, match=msg):
         styler_blank.background_gradient(gmap=gmap, axis=None)._compute()
+
+
+def test_bar_colormap():
+    data = DataFrame([[1, 2], [3, 4]])
+    ctx = data.style.bar(color=mpl.cm.get_cmap("PuBu"), axis=None)._compute().ctx
+    pubu_colors = {
+        (0, 0): "#d0d1e6",
+        (1, 0): "#056faf",
+        (0, 1): "#73a9cf",
+        (1, 1): "#023858",
+    }
+    for k, v in pubu_colors.items():
+        assert v in ctx[k][1][1]
+
+
+def test_bar_color_raises(df):
+    msg = "`colors` must be a matplotlib Colormap if not string or list of strings"
+    with pytest.raises(ValueError, match=msg):
+        df.style.bar(color={"a", "b"}).to_html()
+
+    msg = "`color` must be string, list-like of 2 strings, or matplotlib Colormap"
+    with pytest.raises(ValueError, match=msg):
+        df.style.bar(color=["a", "b", "c"]).to_html()