add new method to styler

Author: quangngd

pandas/io/formats/style.py

@@ -1683,6 +1683,8 @@ def _copy(self, deepcopy: bool = False) -> Styler:
             "_display_funcs",
             "_display_funcs_index",
             "_display_funcs_columns",
+            "_display_funcs_index_names",
+            "_display_funcs_column_names",
             "hidden_rows",
             "hidden_columns",
             "ctx",

pandas/io/formats/style_render.py

@@ -140,9 +140,15 @@ def __init__(
         self._display_funcs_index: DefaultDict[  # maps (row, level) -> format func
             tuple[int, int], Callable[[Any], str]
         ] = defaultdict(lambda: partial(_default_formatter, precision=precision))
+        self._display_funcs_index_names: DefaultDict[  # maps index level -> format func
+            tuple[int, int], Callable[[Any], str]
+        ] = defaultdict(lambda: partial(_default_formatter, precision=precision))
         self._display_funcs_columns: DefaultDict[  # maps (level, col) -> format func
             tuple[int, int], Callable[[Any], str]
         ] = defaultdict(lambda: partial(_default_formatter, precision=precision))
+        self._display_funcs_column_names: DefaultDict[  # maps col level -> format func
+            tuple[int, int], Callable[[Any], str]
+        ] = defaultdict(lambda: partial(_default_formatter, precision=precision))
 
     def _render(
         self,
@@ -460,6 +466,12 @@ def _generate_col_header_row(
         ] * (self.index.nlevels - sum(self.hide_index_) - 1)
 
         name = self.data.columns.names[r]
+
+        is_display = name is not None and not self.hide_column_names
+        value = name if is_display else self.css["blank_value"]
+        display_value = (
+            self._display_funcs_column_names[r](value) if is_display else None
+        )
         column_name = [
             _element(
                 "th",
@@ -468,10 +480,9 @@ def _generate_col_header_row(
                     if name is None
                     else f"{self.css['index_name']} {self.css['level']}{r}"
                 ),
-                name
-                if (name is not None and not self.hide_column_names)
-                else self.css["blank_value"],
+                value,
                 not all(self.hide_index_),
+                display_value=display_value,
             )
         ]
 
@@ -553,6 +564,9 @@ def _generate_index_names_row(
                 f"{self.css['index_name']} {self.css['level']}{c}",
                 self.css["blank_value"] if name is None else name,
                 not self.hide_index_[c],
+                display_value=(
+                    None if name is None else self._display_funcs_index_names[c](name)
+                ),
             )
             for c, name in enumerate(self.data.index.names)
         ]
@@ -1560,6 +1574,139 @@ def alias_(x, value):
 
         return self
 
+    def format_names(
+        self,
+        formatter: ExtFormatter | None = None,
+        axis: Axis = 0,
+        level: Level | list[Level] | None = None,
+        na_rep: str | None = None,
+        precision: int | None = None,
+        decimal: str = ".",
+        thousands: str | None = None,
+        escape: str | None = None,
+        hyperlinks: str | None = None,
+    ) -> StylerRenderer:
+        r"""
+        Format the text display value of index names or column names.
+
+        .. versionadded:: TODO
+
+        Parameters
+        ----------
+        formatter : str, callable, dict or None
+            Object to define how values are displayed. See notes.
+        axis : {0, "index", 1, "columns"}
+            Whether to apply the formatter to the index or column headers.
+        level : int, str, list
+            The level(s) over which to apply the generic formatter.
+        na_rep : str, optional
+            Representation for missing values.
+            If ``na_rep`` is None, no special formatting is applied.
+        precision : int, optional
+            Floating point precision to use for display purposes, if not determined by
+            the specified ``formatter``.
+        decimal : str, default "."
+            Character used as decimal separator for floats, complex and integers.
+        thousands : str, optional, default None
+            Character used as thousands separator for floats, complex and integers.
+        escape : str, optional
+            Use 'html' to replace the characters ``&``, ``<``, ``>``, ``'``, and ``"``
+            in cell display string with HTML-safe sequences.
+            Use 'latex' to replace the characters ``&``, ``%``, ``$``, ``#``, ``_``,
+            ``{``, ``}``, ``~``, ``^``, and ``\`` in the cell display string with
+            LaTeX-safe sequences.
+            Escaping is done before ``formatter``.
+        hyperlinks : {"html", "latex"}, optional
+            Convert string patterns containing https://, http://, ftp:// or www. to
+            HTML <a> tags as clickable URL hyperlinks if "html", or LaTeX \href
+            commands if "latex".
+
+        Returns
+        -------
+        Styler
+
+        See Also
+        --------
+        Styler.format_index: Format the text display value of index labels
+            or column headers.
+
+        Notes
+        -----
+        This method assigns a formatting function, ``formatter``, to each level label
+        in the DataFrame's index or column headers. If ``formatter`` is ``None``,
+        then the default formatter is used.
+        If a callable then that function should take a label value as input and return
+        a displayable representation, such as a string. If ``formatter`` is
+        given as a string this is assumed to be a valid Python format specification
+        and is wrapped to a callable as ``string.format(x)``. If a ``dict`` is given,
+        keys should correspond to MultiIndex level numbers or names, and values should
+        be string or callable, as above.
+
+        The default formatter currently expresses floats and complex numbers with the
+        pandas display precision unless using the ``precision`` argument here. The
+        default formatter does not adjust the representation of missing values unless
+        the ``na_rep`` argument is used.
+
+        The ``level`` argument defines which levels of a MultiIndex to apply the
+        method to. If the ``formatter`` argument is given in dict form but does
+        not include all levels within the level argument then these unspecified levels
+        will have the default formatter applied. Any levels in the formatter dict
+        specifically excluded from the level argument will be ignored.
+
+        When using a ``formatter`` string the dtypes must be compatible, otherwise a
+        `ValueError` will be raised.
+
+        .. warning::
+           `Styler.format_index` is ignored when using the output format
+           `Styler.to_excel`, since Excel and Python have inherrently different
+           formatting structures.
+           However, it is possible to use the `number-format` pseudo CSS attribute
+           to force Excel permissible formatting. See documentation for `Styler.format`.
+        """
+        axis = self.data._get_axis_number(axis)
+        if axis == 0:
+            display_funcs_, obj = self._display_funcs_index_names, self.index
+        else:
+            display_funcs_, obj = self._display_funcs_column_names, self.columns
+        levels_ = refactor_levels(level, obj)
+
+        if all(
+            (
+                formatter is None,
+                level is None,
+                precision is None,
+                decimal == ".",
+                thousands is None,
+                na_rep is None,
+                escape is None,
+                hyperlinks is None,
+            )
+        ):
+            display_funcs_.clear()
+            return self  # clear the formatter / revert to default and avoid looping
+
+        if not isinstance(formatter, dict):
+            formatter = {level: formatter for level in levels_}
+        else:
+            formatter = {
+                obj._get_level_number(level): formatter_
+                for level, formatter_ in formatter.items()
+            }
+
+        for lvl in levels_:
+            format_func = _maybe_wrap_formatter(
+                formatter.get(lvl),
+                na_rep=na_rep,
+                precision=precision,
+                decimal=decimal,
+                thousands=thousands,
+                escape=escape,
+                hyperlinks=hyperlinks,
+            )
+            display_funcs_[lvl] = format_func
+
+        return self
+
 
 def _element(
     html_element: str,
@@ -1571,7 +1718,7 @@ def _element(
     """
     Template to return container with information for a <td></td> or <th></th> element.
     """
-    if "display_value" not in kwargs:
+    if "display_value" not in kwargs or kwargs["display_value"] is None:
         kwargs["display_value"] = value
     return {
         "type": html_element,

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

@@ -77,6 +77,8 @@ def mi_styler_comp(mi_styler):
             columns=mi_styler.columns,
         )
     )
+    mi_styler.format_names(escape="html", axis=0)
+    mi_styler.format_names(escape="html", axis=1)
     return mi_styler