feat: add cols_label_with() method

docs/_quarto.yml

@@ -153,6 +153,7 @@ quartodoc:
         - GT.cols_align
         - GT.cols_width
         - GT.cols_label
+        - GT.cols_label_with
         - GT.cols_move
         - GT.cols_move_to_start
         - GT.cols_move_to_end

great_tables/_boxhead.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import Callable, TYPE_CHECKING
 
 from ._locations import resolve_cols_c
 from ._utils import _assert_list_is_subset
@@ -156,6 +156,75 @@ def cols_label(
     return self._replace(_boxhead=boxhead)
 
 
+def cols_label_with(self: GTSelf, fn: Callable[[str], str], columns: SelectExpr = None) -> GTSelf:
+    """
+    Relabel one or more columns using a function.
+
+    The `cols_label_with()` function allows for modification of column labels through a supplied
+    function. By default, the function will be invoked on all column labels but this can be limited
+    to a subset via the `columns` parameter.
+
+    Parameters
+    ----------
+    fn
+        A function that accepts a column label as input and returns a transformed label as output.
+
+    columns
+        The columns to target. Can either be a single column name or a series of column names
+        provided in a list.
+
+    Returns
+    -------
+    GT
+        The GT object is returned. This is the same object that the method is called on so that we
+        can facilitate method chaining.
+
+    Notes
+    -----
+    GT always selects columns using their name in the underlying data. This means that a column's
+    label is purely for final presentation.
+
+    Examples
+    --------
+    Let's use a subset of the `sp500` dataset to create a gt table.
+    ```{python}
+    from great_tables import GT, md
+    from great_tables.data import sp500
+
+    gt = GT(sp500.head())
+    gt
+    ```
+
+    We can pass `str.upper()` to the `columns` parameter to convert all column labels to uppercase.
+    ```{python}
+    gt.cols_label_with(str.upper)
+    ```
+
+    One useful use case is using `md()`, provided by **Great Tables**, to format column labels.
+    For example, the following code demonstrates how to make the `date` and `adj_close` column labels
+    bold using markdown syntax.
+    ```{python}
+    gt.cols_label_with(lambda x: md(f"**{x}**"), columns=["date", "adj_close"])
+    ```
+
+    """
+    # Get the full list of column names for the data
+    column_names = self._boxhead._get_columns()
+
+    if isinstance(columns, str):
+        columns = [columns]
+        _assert_list_is_subset(columns, set_list=column_names)
+    elif columns is None:
+        columns = column_names
+
+    sel_cols = resolve_cols_c(data=self, expr=columns)
+
+    new_col_labels = {col: fn(col) for col in sel_cols}
+    boxhead = self._boxhead._set_column_labels(new_col_labels)
+
+    return self._replace(_boxhead=boxhead)
+
+
 def cols_align(self: GTSelf, align: str = "left", columns: SelectExpr = None) -> GTSelf:
     """
     Set the alignment of one or more columns.

great_tables/gt.py

@@ -6,7 +6,7 @@
 
 # Main gt imports ----
 from ._body import body_reassemble
-from ._boxhead import cols_align, cols_label
+from ._boxhead import cols_align, cols_label, cols_label_with
 from ._data_color import data_color
 from ._export import as_latex, as_raw_html, save, show, write_raw_html
 from ._formats import (
@@ -253,6 +253,7 @@ def __init__(
     cols_align = cols_align
     cols_width = cols_width
     cols_label = cols_label
+    cols_label_with = cols_label_with
     cols_move = cols_move
     cols_move_to_start = cols_move_to_start
     cols_move_to_end = cols_move_to_end

tests/test__boxhead.py

@@ -56,6 +56,33 @@ def test_cols_label_return_self_if_no_kwargs():
     assert isinstance(unmodified_table, gt.GT)
 
 
+def test_cols_label_with_relabel_columns():
+    # Create a table with default column labels
+    df = pd.DataFrame({"A": [1, 2, 3], "B": [4, 5, 6]})
+    table = gt.GT(df)
+
+    # Relabel the columns
+    modified_table = table.cols_label_with(str.lower)
+
+    # Check that the column labels have been updated
+    assert modified_table._boxhead._get_column_labels() == ["a", "b"]
+
+
+def test_cols_label_with_relabel_columns_with_markdown():
+    # Create a table with default column labels
+    df = pd.DataFrame({"A": [1, 2, 3], "B": [4, 5, 6]})
+    table = gt.GT(df)
+
+    # Relabel a column with a Markdown formatted label
+    modified_table = table.cols_label_with(lambda x: gt.md(f"**{x}**"), columns="A")
+
+    # Check that the column label has been updated with Markdown formatting
+    modified_column_labels = modified_table._boxhead._get_column_labels()
+
+    assert modified_column_labels[0].text == "**A**"
+    assert modified_column_labels[1] == "B"
+
+
 def test_cols_align_default():
     df = pd.DataFrame({"A": [1, 2, 3], "B": [4, 5, 6]})
     table = gt.GT(df)