Backport PR #42072: ENH: Styler.set_sticky for maintaining index and column headers in HTML frame (#42096)

Co-authored-by: attack68 <[email protected]>

Author: meeseeksmachine

doc/source/reference/style.rst

@@ -42,6 +42,7 @@ Style application
    Styler.set_table_attributes
    Styler.set_tooltips
    Styler.set_caption
+   Styler.set_sticky
    Styler.set_properties
    Styler.set_uuid
    Styler.clear

doc/source/user_guide/style.ipynb

@@ -1405,7 +1405,26 @@
    "source": [
     "### Sticky Headers\n",
     "\n",
-    "If you display a large matrix or DataFrame in a notebook, but you want to always see the column and row headers you can use the following CSS to make them stick. We might make this into an API function later."
+    "If you display a large matrix or DataFrame in a notebook, but you want to always see the column and row headers you can use the [.set_sticky][sticky] method which manipulates the table styles CSS.\n",
+    "\n",
+    "[sticky]: ../reference/api/pandas.io.formats.style.Styler.set_sticky.rst"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "bigdf = pd.DataFrame(np.random.randn(16, 100))\n",
+    "bigdf.style.set_sticky(axis=\"index\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "It is also possible to stick MultiIndexes and even only specific levels."
    ]
   },
   {
@@ -1414,11 +1433,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "bigdf = pd.DataFrame(np.random.randn(15, 100))\n",
-    "bigdf.style.set_table_styles([\n",
-    "    {'selector': 'thead th', 'props': 'position: sticky; top:0; background-color:salmon;'},\n",
-    "    {'selector': 'tbody th', 'props': 'position: sticky; left:0; background-color:lightgreen;'}  \n",
-    "])"
+    "bigdf.index = pd.MultiIndex.from_product([[\"A\",\"B\"],[0,1],[0,1,2,3]])\n",
+    "bigdf.style.set_sticky(axis=\"index\", pixel_size=18, levels=[1,2])"
    ]
   },
   {

doc/source/whatsnew/v1.3.0.rst

@@ -138,6 +138,7 @@ which has been revised and improved (:issue:`39720`, :issue:`39317`, :issue:`404
  - Added the option ``styler.render.max_elements`` to avoid browser overload when styling large DataFrames (:issue:`40712`)
  - Added the method :meth:`.Styler.to_latex` (:issue:`21673`), which also allows some limited CSS conversion (:issue:`40731`)
  - Added the method :meth:`.Styler.to_html` (:issue:`13379`)
+ - Added the method :meth:`.Styler.set_sticky` to make index and column headers permanently visible in scrolling HTML frames (:issue:`29072`)
 
 .. _whatsnew_130.enhancements.dataframe_honors_copy_with_dict:
 

pandas/io/formats/style.py

@@ -1414,6 +1414,71 @@ def set_caption(self, caption: str | tuple) -> Styler:
         self.caption = caption
         return self
 
+    def set_sticky(
+        self,
+        axis: Axis = 0,
+        pixel_size: int | None = None,
+        levels: list[int] | None = None,
+    ) -> Styler:
+        """
+        Add CSS to permanently display the index or column headers in a scrolling frame.
+
+        Parameters
+        ----------
+        axis : {0 or 'index', 1 or 'columns', None}, default 0
+            Whether to make the index or column headers sticky.
+        pixel_size : int, optional
+            Required to configure the width of index cells or the height of column
+            header cells when sticking a MultiIndex. Defaults to 75 and 25 respectively.
+        levels : list of int
+            If ``axis`` is a MultiIndex the specific levels to stick. If ``None`` will
+            stick all levels.
+
+        Returns
+        -------
+        self : Styler
+        """
+        if axis in [0, "index"]:
+            axis, obj, tag, pos = 0, self.data.index, "tbody", "left"
+            pixel_size = 75 if not pixel_size else pixel_size
+        elif axis in [1, "columns"]:
+            axis, obj, tag, pos = 1, self.data.columns, "thead", "top"
+            pixel_size = 25 if not pixel_size else pixel_size
+        else:
+            raise ValueError("`axis` must be one of {0, 1, 'index', 'columns'}")
+
+        if not isinstance(obj, pd.MultiIndex):
+            return self.set_table_styles(
+                [
+                    {
+                        "selector": f"{tag} th",
+                        "props": f"position:sticky; {pos}:0px; background-color:white;",
+                    }
+                ],
+                overwrite=False,
+            )
+        else:
+            range_idx = list(range(obj.nlevels))
+
+        levels = sorted(levels) if levels else range_idx
+        for i, level in enumerate(levels):
+            self.set_table_styles(
+                [
+                    {
+                        "selector": f"{tag} th.level{level}",
+                        "props": f"position: sticky; "
+                        f"{pos}: {i * pixel_size}px; "
+                        f"{f'height: {pixel_size}px; ' if axis == 1 else ''}"
+                        f"{f'min-width: {pixel_size}px; ' if axis == 0 else ''}"
+                        f"{f'max-width: {pixel_size}px; ' if axis == 0 else ''}"
+                        f"background-color: white;",
+                    }
+                ],
+                overwrite=False,
+            )
+
+        return self
+
     def set_table_styles(
         self,
         table_styles: dict[Any, CSSStyles] | CSSStyles,

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

@@ -1,8 +1,12 @@
 from textwrap import dedent
 
+import numpy as np
 import pytest
 
-from pandas import DataFrame
+from pandas import (
+    DataFrame,
+    MultiIndex,
+)
 
 jinja2 = pytest.importorskip("jinja2")
 from pandas.io.formats.style import Styler
@@ -16,6 +20,12 @@ def styler():
     return Styler(DataFrame([[2.61], [2.69]], index=["a", "b"], columns=["A"]))
 
 
+@pytest.fixture
+def styler_mi():
+    midx = MultiIndex.from_product([["a", "b"], ["c", "d"]])
+    return Styler(DataFrame(np.arange(16).reshape(4, 4), index=midx, columns=midx))
+
+
 @pytest.fixture
 def tpl_style():
     return env.get_template("html_style.tpl")
@@ -236,3 +246,146 @@ def test_from_custom_template(tmpdir):
 def test_caption_as_sequence(styler):
     styler.set_caption(("full cap", "short cap"))
     assert "<caption>full cap</caption>" in styler.render()
+
+
+@pytest.mark.parametrize("index", [False, True])
+@pytest.mark.parametrize("columns", [False, True])
+def test_sticky_basic(styler, index, columns):
+    if index:
+        styler.set_sticky(axis=0)
+    if columns:
+        styler.set_sticky(axis=1)
+
+    res = styler.set_uuid("").to_html()
+    cs1 = "tbody th {\n  position: sticky;\n  left: 0px;\n  background-color: white;\n}"
+    assert (cs1 in res) is index
+    cs2 = "thead th {\n  position: sticky;\n  top: 0px;\n  background-color: white;\n}"
+    assert (cs2 in res) is columns
+
+
+@pytest.mark.parametrize("index", [False, True])
+@pytest.mark.parametrize("columns", [False, True])
+def test_sticky_mi(styler_mi, index, columns):
+    if index:
+        styler_mi.set_sticky(axis=0)
+    if columns:
+        styler_mi.set_sticky(axis=1)
+
+    res = styler_mi.set_uuid("").to_html()
+    assert (
+        (
+            dedent(
+                """\
+        #T_ tbody th.level0 {
+          position: sticky;
+          left: 0px;
+          min-width: 75px;
+          max-width: 75px;
+          background-color: white;
+        }
+        """
+            )
+            in res
+        )
+        is index
+    )
+    assert (
+        (
+            dedent(
+                """\
+        #T_ tbody th.level1 {
+          position: sticky;
+          left: 75px;
+          min-width: 75px;
+          max-width: 75px;
+          background-color: white;
+        }
+        """
+            )
+            in res
+        )
+        is index
+    )
+    assert (
+        (
+            dedent(
+                """\
+        #T_ thead th.level0 {
+          position: sticky;
+          top: 0px;
+          height: 25px;
+          background-color: white;
+        }
+        """
+            )
+            in res
+        )
+        is columns
+    )
+    assert (
+        (
+            dedent(
+                """\
+        #T_ thead th.level1 {
+          position: sticky;
+          top: 25px;
+          height: 25px;
+          background-color: white;
+        }
+        """
+            )
+            in res
+        )
+        is columns
+    )
+
+
+@pytest.mark.parametrize("index", [False, True])
+@pytest.mark.parametrize("columns", [False, True])
+def test_sticky_levels(styler_mi, index, columns):
+    if index:
+        styler_mi.set_sticky(axis=0, levels=[1])
+    if columns:
+        styler_mi.set_sticky(axis=1, levels=[1])
+
+    res = styler_mi.set_uuid("").to_html()
+    assert "#T_ tbody th.level0 {" not in res
+    assert "#T_ thead th.level0 {" not in res
+    assert (
+        (
+            dedent(
+                """\
+        #T_ tbody th.level1 {
+          position: sticky;
+          left: 0px;
+          min-width: 75px;
+          max-width: 75px;
+          background-color: white;
+        }
+        """
+            )
+            in res
+        )
+        is index
+    )
+    assert (
+        (
+            dedent(
+                """\
+        #T_ thead th.level1 {
+          position: sticky;
+          top: 0px;
+          height: 25px;
+          background-color: white;
+        }
+        """
+            )
+            in res
+        )
+        is columns
+    )
+
+
+def test_sticky_raises(styler):
+    with pytest.raises(ValueError, match="`axis` must be"):
+        styler.set_sticky(axis="bad")