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,9 +1,9 @@
 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
+from ._utils import _assert_list_is_subset, _handle_units_syntax
 from ._tbl_data import SelectExpr
 from ._text import BaseText
 
@@ -114,8 +114,6 @@ def cols_label(
     )
     ```
     """
-    from great_tables._helpers import UnitStr
-
     cases = cases if cases is not None else {}
     new_cases = cases | kwargs
 
@@ -132,24 +130,80 @@ def cols_label(
     _assert_list_is_subset(mod_columns, set_list=column_names)
 
     # Handle units syntax in labels (e.g., "Density ({{ppl / mi^2}})")
-    new_kwargs: dict[str, UnitStr | str | BaseText] = {}
+    new_kwargs = _handle_units_syntax(new_cases)
+
+    boxhead = self._boxhead._set_column_labels(new_kwargs)
+
+    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.
 
-    for k, v in new_cases.items():
-        if isinstance(v, str):
-            unitstr_v = UnitStr.from_str(v)
+    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
 
-            if len(unitstr_v.units_str) == 1 and isinstance(unitstr_v.units_str[0], str):
-                new_kwargs[k] = unitstr_v.units_str[0]
-            else:
-                new_kwargs[k] = unitstr_v
+    gt = GT(sp500.head())
+    gt
+    ```
 
-        elif isinstance(v, BaseText):
-            new_kwargs[k] = v
+    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"])
+    ```
 
-        else:
-            raise ValueError(
-                "Column labels must be strings or BaseText objects. Use `md()` or `html()` for formatting."
-            )
+    """
+    # 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_cases = {col: fn(col) for col in sel_cols}
+
+    # Handle units syntax in labels (e.g., "Density ({{ppl / mi^2}})")
+    new_kwargs = _handle_units_syntax(new_cases)
 
     boxhead = self._boxhead._set_column_labels(new_kwargs)
 

great_tables/_utils.py

@@ -7,6 +7,7 @@
 from types import ModuleType
 from typing import TYPE_CHECKING, Any, Iterable, Iterator
 
+from ._helpers import UnitStr
 from ._tbl_data import _get_cell, _set_cell, get_column_names, n_rows
 from ._text import BaseText, _process_text
 
@@ -285,3 +286,26 @@ def _get_visible_cells(data: TblData) -> list[tuple[str, int]]:
 
 def is_valid_http_schema(url: str) -> bool:
     return url.startswith("http://") or url.startswith("https://")
+
+
+def _handle_units_syntax(cases: dict[str, str | BaseText]) -> dict[str, UnitStr | str | BaseText]:
+    # Handle units syntax in labels (e.g., "Density ({{ppl / mi^2}})")
+    kwargs: dict[str, UnitStr | str | BaseText] = {}
+
+    for k, v in cases.items():
+        if isinstance(v, str):
+            unitstr_v = UnitStr.from_str(v)
+
+            if len(unitstr_v.units_str) == 1 and isinstance(unitstr_v.units_str[0], str):
+                kwargs[k] = unitstr_v.units_str[0]
+            else:
+                kwargs[k] = unitstr_v
+
+        elif isinstance(v, BaseText):
+            kwargs[k] = v
+
+        else:
+            raise ValueError(
+                "Column labels must be strings or BaseText objects. Use `md()` or `html()` for formatting."
+            )
+    return kwargs

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)

tests/test_utils.py

@@ -10,6 +10,7 @@
     _assert_str_list,
     _assert_str_scalar,
     _collapse_list_elements,
+    _handle_units_syntax,
     _insert_into_list,
     _match_arg,
     _migrate_unformatted_to_output,
@@ -224,3 +225,21 @@ def test_migrate_unformatted_to_output_html():
 )
 def test_is_valid_http_schema(url: str):
     assert is_valid_http_schema(url)
+
+
+def test_handle_units_syntax():
+    from great_tables._text import BaseText, Text
+
+    new_kwargs = _handle_units_syntax({"column_label_1": "abc", "column_label_2": Text(text="xyz")})
+
+    assert all(isinstance(v, (str, BaseText)) for v in new_kwargs.values())
+
+
+def test_handle_units_syntax_raises():
+    with pytest.raises(ValueError) as exc_info:
+        _handle_units_syntax({"column_label": 123})
+
+    assert (
+        "Column labels must be strings or BaseText objects. Use `md()` or `html()` for formatting."
+        in exc_info.value.args[0]
+    )