ENH: Styler.apply_index and Styler.applymap_index for conditional formatting of column/index headers

doc/source/reference/style.rst

@@ -36,6 +36,8 @@ Style application
 
    Styler.apply
    Styler.applymap
+   Styler.apply_header
+   Styler.applymap_header
    Styler.format
    Styler.hide_index
    Styler.hide_columns

doc/source/whatsnew/v1.4.0.rst

@@ -33,6 +33,7 @@ Other enhancements
 - :meth:`Series.sample`, :meth:`DataFrame.sample`, and :meth:`.GroupBy.sample` now accept a ``np.random.Generator`` as input to ``random_state``. A generator will be more performant, especially with ``replace=False`` (:issue:`38100`)
 -  Additional options added to :meth:`.Styler.bar` to control alignment and display, with keyword only arguments (:issue:`26070`, :issue:`36419`)
 - :meth:`Series.ewm`, :meth:`DataFrame.ewm`, now support a ``method`` argument with a ``'table'`` option that performs the windowing operation over an entire :class:`DataFrame`. See :ref:`Window Overview <window.overview>` for performance and functional benefits (:issue:`42273`)
+- :meth:`.Styler.apply_header` and :meth:`.Styler.applymap_header` added to allow conditional styling of index and column header values (:issue:`41893`)
 -
 
 .. ---------------------------------------------------------------------------

pandas/io/formats/style.py

@@ -964,6 +964,32 @@ def _update_ctx(self, attrs: DataFrame) -> None:
                 i, j = self.index.get_loc(rn), self.columns.get_loc(cn)
                 self.ctx[(i, j)].extend(css_list)
 
+    def _update_ctx_header(self, attrs: DataFrame, axis: str) -> None:
+        """
+        Update the state of the ``Styler`` for header cells.
+
+        Collects a mapping of {index_label: [('<property>', '<value>'), ..]}.
+
+        Parameters
+        ----------
+        attrs : Series
+            Should contain strings of '<property>: <value>;<prop2>: <val2>', and an
+            integer index.
+            Whitespace shouldn't matter and the final trailing ';' shouldn't
+            matter.
+        axis : str
+            Identifies whether the ctx object being updated is the index or columns
+        """
+        for j in attrs.columns:
+            for i, c in attrs[[j]].itertuples():
+                if not c:
+                    continue
+                css_list = maybe_convert_css_to_tuples(c)
+                if axis == "index":
+                    self.ctx_index[(i, j)].extend(css_list)
+                else:
+                    self.ctx_columns[(j, i)].extend(css_list)
+
     def _copy(self, deepcopy: bool = False) -> Styler:
         """
         Copies a Styler, allowing for deepcopy or shallow copy
@@ -1003,6 +1029,8 @@ def _copy(self, deepcopy: bool = False) -> Styler:
             "hidden_rows",
             "hidden_columns",
             "ctx",
+            "ctx_index",
+            "ctx_columns",
             "cell_context",
             "_todo",
             "table_styles",
@@ -1124,6 +1152,8 @@ def apply(
 
         See Also
         --------
+        Styler.applymap_header: Apply a CSS-styling function to headers elementwise.
+        Styler.apply_header: Apply a CSS-styling function to headers level-wise.
         Styler.applymap: Apply a CSS-styling function elementwise.
 
         Notes
@@ -1160,6 +1190,151 @@ def apply(
         )
         return self
 
+    def _apply_header(
+        self,
+        func: Callable[..., Styler],
+        axis: int | str = 0,
+        levels: list[int] | int | None = None,
+        method: str = "apply",
+        **kwargs,
+    ) -> Styler:
+        if axis in [0, "index"]:
+            obj, axis = self.index, "index"
+        elif axis in [1, "columns"]:
+            obj, axis = self.columns, "columns"
+        else:
+            raise ValueError(
+                f"`axis` must be one of 0, 1, 'index', 'columns', got {axis}"
+            )
+
+        if isinstance(obj, pd.MultiIndex) and levels is not None:
+            levels = [levels] if isinstance(levels, int) else levels
+            data = DataFrame(obj.to_list()).loc[:, levels]
+        else:
+            data = DataFrame(obj.to_list())
+
+        if method == "apply":
+            result = data.apply(func, axis=0, **kwargs)
+        elif method == "applymap":
+            result = data.applymap(func, **kwargs)
+
+        self._update_ctx_header(result, axis)
+        return self
+
+    @doc(
+        this="apply",
+        wise="level-wise",
+        alt="applymap",
+        altwise="elementwise",
+        func="take a Series and return a string array of the same length",
+        axis='{0, 1, "index", "columns"}',
+        input_note="the index as a Series, if an Index, or a level of a MultiIndex",
+        output_note="an identically sized array of CSS styles as strings",
+        var="s",
+        ret='np.where(s == "B", "background-color: yellow;", "")',
+        ret2='["background-color: yellow;" if "x" in v else "" for v in s]',
+    )
+    def apply_header(
+        self,
+        func: Callable[..., Styler],
+        axis: int | str = 0,
+        levels: list[int] | int | None = None,
+        **kwargs,
+    ) -> Styler:
+        """
+        Apply a CSS-styling function to the index or column headers, {wise}.
+
+        Updates the HTML representation with the result.
+
+        .. versionadded:: 1.4.0
+
+        Parameters
+        ----------
+        func : function
+            ``func`` should {func}.
+        axis : {axis}
+            The headers over which to apply the function.
+        levels : int, list of ints, optional
+            If index is MultiIndex the level(s) over which to apply the function.
+        **kwargs : dict
+            Pass along to ``func``.
+
+        Returns
+        -------
+        self : Styler
+
+        See Also
+        --------
+        Styler.{alt}_header: Apply a CSS-styling function to headers {altwise}.
+        Styler.apply: Apply a CSS-styling function column-wise, row-wise, or table-wise.
+        Styler.applymap: Apply a CSS-styling function elementwise.
+
+        Notes
+        -----
+        Each input to ``func`` will be {input_note}. The output of ``func`` should be
+        {output_note}, in the format 'attribute: value; attribute2: value2; ...'
+        or, if nothing is to be applied to that element, an empty string or ``None``.
+
+        Examples
+        --------
+        Basic usage to conditionally highlight values in the index.
+
+        >>> df = pd.DataFrame([[1,2], [3,4]], index=["A", "B"])
+        >>> def color_b(s):
+        ...     return {ret}
+        >>> df.style.{this}_header(color_b)
+
+        .. figure:: ../../_static/style/appmaphead1.png
+
+        Selectively applying to specific levels of MultiIndex columns.
+
+        >>> midx = pd.MultiIndex.from_product([['ix', 'jy'], [0, 1], ['x3', 'z4']])
+        >>> df = pd.DataFrame([np.arange(8)], columns=midx)
+        >>> def highlight_x({var}):
+        ...     return {ret2}
+        >>> df.style.{this}_header(highlight_x, axis="columns", levels=[0, 2])
+
+        .. figure:: ../../_static/style/appmaphead2.png
+        """
+        self._todo.append(
+            (
+                lambda instance: getattr(instance, "_apply_header"),
+                (func, axis, levels, "apply"),
+                kwargs,
+            )
+        )
+        return self
+
+    @doc(
+        apply_header,
+        this="applymap",
+        wise="elementwise",
+        alt="apply",
+        altwise="level-wise",
+        func="take a scalar and return a string",
+        axis='{0, 1, "index", "columns"}',
+        input_note="an index value, if an Index, or a level value of a MultiIndex",
+        output_note="CSS styles as a string",
+        var="v",
+        ret='"background-color: yellow;" if v == "B" else None',
+        ret2='"background-color: yellow;" if "x" in v else None',
+    )
+    def applymap_header(
+        self,
+        func: Callable[..., Styler],
+        axis: int | str = 0,
+        levels: list[int] | int | None = None,
+        **kwargs,
+    ) -> Styler:
+        self._todo.append(
+            (
+                lambda instance: getattr(instance, "_apply_header"),
+                (func, axis, levels, "applymap"),
+                kwargs,
+            )
+        )
+        return self
+
     def _applymap(
         self, func: Callable, subset: Subset | None = None, **kwargs
     ) -> Styler:
@@ -1182,7 +1357,7 @@ def applymap(
         Parameters
         ----------
         func : function
-            ``func`` should take a scalar and return a scalar.
+            ``func`` should take a scalar and return a string.
         subset : label, array-like, IndexSlice, optional
             A valid 2d input to `DataFrame.loc[<subset>]`, or, in the case of a 1d input
             or single key, to `DataFrame.loc[:, <subset>]` where the columns are
@@ -1196,6 +1371,8 @@ def applymap(
 
         See Also
         --------
+        Styler.applymap_header: Apply a CSS-styling function to headers elementwise.
+        Styler.apply_header: Apply a CSS-styling function to headers level-wise.
         Styler.apply: Apply a CSS-styling function column-wise, row-wise, or table-wise.
 
         Notes

pandas/io/formats/style_render.py

@@ -102,6 +102,8 @@ def __init__(
         self.hidden_rows: Sequence[int] = []  # sequence for specific hidden rows/cols
         self.hidden_columns: Sequence[int] = []
         self.ctx: DefaultDict[tuple[int, int], CSSList] = defaultdict(list)
+        self.ctx_index: DefaultDict[tuple[int, int], CSSList] = defaultdict(list)
+        self.ctx_columns: DefaultDict[tuple[int, int], CSSList] = defaultdict(list)
         self.cell_context: DefaultDict[tuple[int, int], str] = defaultdict(str)
         self._todo: list[tuple[Callable, tuple, dict]] = []
         self.tooltips: Tooltips | None = None
@@ -152,6 +154,8 @@ def _compute(self):
         (application method, *args, **kwargs)
         """
         self.ctx.clear()
+        self.ctx_index.clear()
+        self.ctx_columns.clear()
         r = self
         for func, args, kwargs in self._todo:
             r = func(self)(*args, **kwargs)
@@ -201,6 +205,9 @@ def _translate(self, sparse_index: bool, sparse_cols: bool, blank: str = " 
             len(self.data.index), len(self.data.columns), max_elements
         )
 
+        self.cellstyle_map_columns: DefaultDict[
+            tuple[CSSPair, ...], list[str]
+        ] = defaultdict(list)
         head = self._translate_header(
             BLANK_CLASS,
             BLANK_VALUE,
@@ -215,6 +222,9 @@ def _translate(self, sparse_index: bool, sparse_cols: bool, blank: str = " 
         self.cellstyle_map: DefaultDict[tuple[CSSPair, ...], list[str]] = defaultdict(
             list
         )
+        self.cellstyle_map_index: DefaultDict[
+            tuple[CSSPair, ...], list[str]
+        ] = defaultdict(list)
         body = self._translate_body(
             DATA_CLASS,
             ROW_HEADING_CLASS,
@@ -226,11 +236,17 @@ def _translate(self, sparse_index: bool, sparse_cols: bool, blank: str = " 
         )
         d.update({"body": body})
 
-        cellstyle: list[dict[str, CSSList | list[str]]] = [
-            {"props": list(props), "selectors": selectors}
-            for props, selectors in self.cellstyle_map.items()
-        ]
-        d.update({"cellstyle": cellstyle})
+        ctx_maps = {
+            "cellstyle": "cellstyle_map",
+            "cellstyle_index": "cellstyle_map_index",
+            "cellstyle_columns": "cellstyle_map_columns",
+        }  # add the cell_ids styles map to the render dictionary in right format
+        for k, attr in ctx_maps.items():
+            map = [
+                {"props": list(props), "selectors": selectors}
+                for props, selectors in getattr(self, attr).items()
+            ]
+            d.update({k: map})
 
         table_attr = self.table_attributes
         use_mathjax = get_option("display.html.use_mathjax")
@@ -320,8 +336,9 @@ def _translate_header(
                 ]
 
                 if clabels:
-                    column_headers = [
-                        _element(
+                    column_headers = []
+                    for c, value in enumerate(clabels[r]):
+                        header_element = _element(
                             "th",
                             f"{col_heading_class} level{r} col{c}",
                             value,
@@ -332,8 +349,16 @@ def _translate_header(
                                 else ""
                             ),
                         )
-                        for c, value in enumerate(clabels[r])
-                    ]
+
+                        if self.cell_ids:
+                            header_element["id"] = f"level{r}_col{c}"
+                        if (r, c) in self.ctx_columns and self.ctx_columns[r, c]:
+                            header_element["id"] = f"level{r}_col{c}"
+                            self.cellstyle_map_columns[
+                                tuple(self.ctx_columns[r, c])
+                            ].append(f"level{r}_col{c}")
+
+                        column_headers.append(header_element)
 
                     if len(self.data.columns) > max_cols:
                         # add an extra column with `...` value to indicate trimming
@@ -466,21 +491,30 @@ def _translate_body(
                 body.append(index_headers + data)
                 break
 
-            index_headers = [
-                _element(
+            index_headers = []
+            for c, value in enumerate(rlabels[r]):
+                header_element = _element(
                     "th",
                     f"{row_heading_class} level{c} row{r}",
                     value,
                     (_is_visible(r, c, idx_lengths) and not self.hide_index_),
-                    id=f"level{c}_row{r}",
                     attributes=(
                         f'rowspan="{idx_lengths.get((c, r), 0)}"'
                         if idx_lengths.get((c, r), 0) > 1
                         else ""
                     ),
                 )
-                for c, value in enumerate(rlabels[r])
-            ]
+
+                if self.cell_ids:
+                    header_element["id"] = f"level{c}_row{r}"  # id is specified
+                if (r, c) in self.ctx_index and self.ctx_index[r, c]:
+                    # always add id if a style is specified
+                    header_element["id"] = f"level{c}_row{r}"
+                    self.cellstyle_map_index[tuple(self.ctx_index[r, c])].append(
+                        f"level{c}_row{r}"
+                    )
+
+                index_headers.append(header_element)
 
             data = []
             for c, value in enumerate(row_tup[1:]):
@@ -510,13 +544,12 @@ def _translate_body(
                     display_value=self._display_funcs[(r, c)](value),
                 )
 
-                # only add an id if the cell has a style
-                if self.cell_ids or (r, c) in self.ctx:
+                if self.cell_ids:
                     data_element["id"] = f"row{r}_col{c}"
-                    if (r, c) in self.ctx and self.ctx[r, c]:  # only add  if non-empty
-                        self.cellstyle_map[tuple(self.ctx[r, c])].append(
-                            f"row{r}_col{c}"
-                        )
+                if (r, c) in self.ctx and self.ctx[r, c]:
+                    # always add id if needed due to specified style
+                    data_element["id"] = f"row{r}_col{c}"
+                    self.cellstyle_map[tuple(self.ctx[r, c])].append(f"row{r}_col{c}")
 
                 data.append(data_element)