Merge pull request #83 from posit-dev/feat-style-from-column

feat: add style.from_column, implement for loc.body

docs/_quarto.yml

@@ -119,6 +119,7 @@ quartodoc:
       contents:
         - md
         - html
+        - from_column
         #- px
         #- pct
     - title: Value formatting functions

great_tables/__init__.py

@@ -12,6 +12,7 @@
 from . import vals
 from . import loc
 from . import style
+from ._styles import FromColumn as from_column
 from ._helpers import letters, LETTERS, px, pct, md, html, random_id
 from .data import exibble
 
@@ -26,6 +27,7 @@
     "md",
     "html",
     "random_id",
+    "from_column",
     "vals",
     "loc",
     "style",

great_tables/_locations.py

@@ -366,6 +366,11 @@ def set_style(loc: Loc, data: GTData, style: List[str]) -> GTData:
 
 @set_style.register
 def _(loc: LocTitle, data: GTData, style: List[CellStyle]) -> GTData:
+    # validate ----
+    for entry in style:
+        entry._raise_if_requires_data(loc)
+
+    # set ----
     if loc.groups == "title":
         info = StyleInfo(locname="title", locnum=1, styles=style)
     elif loc.groups == "subtitle":
@@ -378,12 +383,16 @@ def _(loc: LocTitle, data: GTData, style: List[CellStyle]) -> GTData:
 
 @set_style.register
 def _(loc: LocBody, data: GTData, style: List[CellStyle]) -> GTData:
-    positions = resolve(loc, data)
+    positions: List[CellPos] = resolve(loc, data)
+
+    # evaluate any column expressions in styles
+    style_ready = [entry._evaluate_expressions(data._tbl_data) for entry in style]
 
     all_info: list[StyleInfo] = []
     for col_pos in positions:
+        row_styles = [entry._from_row(data._tbl_data, col_pos.row) for entry in style_ready]
         crnt_info = StyleInfo(
-            locname="data", locnum=5, colname=col_pos.colname, rownum=col_pos.row, styles=style
+            locname="data", locnum=5, colname=col_pos.colname, rownum=col_pos.row, styles=row_styles
         )
         all_info.append(crnt_info)
 

great_tables/_styles.py

@@ -1,21 +1,121 @@
 from __future__ import annotations
 
-from dataclasses import dataclass
-from typing import Literal, List
+from dataclasses import dataclass, fields, replace
+from typing import TYPE_CHECKING, Any, Callable, Literal, List, Union
+from typing_extensions import Self, TypeAlias
+
+from ._tbl_data import TblData, _get_cell, PlExpr, eval_transform
+
+
+if TYPE_CHECKING:
+    from ._locations import Loc
 
 
 # Cell Styles ==========================================================================
 # TODO: stubbed out the styles in helpers.R as dataclasses while I was reading it,
 # but have no worked on any runtime validation, etc..
+ColumnExpr: TypeAlias = Union["FromColumn", PlExpr, "FromValues"]
+
+
+@dataclass
+class FromColumn:
+    """Specify that a style value should be fetched from a column in the data.
+
+    Examples
+    --------
+
+    ```{python}
+    import pandas as pd
+    from great_tables import GT, exibble, from_column, loc, style
+
+    df = pd.DataFrame({"x": [1, 2], "color": ["red", "blue"]})
+
+    gt = GT(df)
+    gt.tab_style(
+        style = style.text(color = from_column("color")),
+        locations = loc.body(columns = ["x"])
+    )
+    ```
+
+    If you are using polars, you can just pass polars expressions in directly:
+
+    ```{python}
+    import polars as pl
+
+    gt_polars = GT(pl.from_pandas(df))
+    gt_polars.tab_style(
+        style = style.text(color = pl.col("color")),    # <-- polars expression
+        locations = loc.body(columns = ["x"])
+    )
+    ```
+    """
+
+    column: str
+    # TODO: na_value currently unused
+    na_value: Any | None = None
+    fn: Callable[[Any], Any] | None = None
+
+
+@dataclass
+class FromValues:
+    values: list[Any]
+    expr: PlExpr | None = None
 
 
 # TODO: what goes into CellStyle?
+@dataclass
 class CellStyle:
     """A style specification."""
 
     def _to_html_style(self) -> str:
         raise NotImplementedError
 
+    def _evaluate_expressions(self, data: TblData) -> Self:
+        new_fields: dict[str, FromValues] = {}
+        for field in fields(self):
+            attr = getattr(self, field.name)
+            if isinstance(attr, PlExpr):
+                col_res = eval_transform(data, attr)
+                new_fields[field.name] = FromValues(expr=attr, values=col_res)
+
+        if not new_fields:
+            return self
+
+        return replace(self, **new_fields)
+
+    def _from_row(self, data: TblData, row: int) -> Self:
+        """Return a new object with FromColumn replaced with values from row.
+
+        Note that if no FromColumn fields are present, this returns the original object.
+        """
+
+        new_fields: dict[str, Any] = {}
+        for field in fields(self):
+            attr = getattr(self, field.name)
+            if isinstance(attr, FromColumn):
+                # TODO: could validate that the value fetched from data is allowed.
+                # e.g. that color is a string, etc..
+                val = _get_cell(data, row, attr.column)
+
+                new_fields[field.name] = attr.fn(val) if attr.fn is not None else val
+            elif isinstance(attr, FromValues):
+                new_fields[field.name] = attr.values[row]
+
+        if not new_fields:
+            return self
+
+        return replace(self, **new_fields)
+
+    def _raise_if_requires_data(self, loc: Loc):
+        for field in fields(self):
+            attr = getattr(self, field.name)
+            if isinstance(attr, FromColumn):
+                raise TypeError(
+                    f"Location type {type(loc)} cannot use FromColumn."
+                    f"\n\nStyle type: {type(self)}"
+                    f"\nField with FromColumn: {field.name}"
+                )
+
 
 @dataclass
 class CellStyleText(CellStyle):
@@ -75,13 +175,13 @@ class CellStyleText(CellStyle):
         properties.
     """
 
-    color: str | None = None
-    font: str | None = None
-    size: str | None = None
-    align: Literal["center", "left", "right", "justify"] | None = None
-    v_align: Literal["middle", "top", "bottom"] | None = None
-    style: Literal["normal", "italic", "oblique"] | None = None
-    weight: Literal["normal", "bold", "bolder", "lighter"] | None = None
+    color: str | ColumnExpr | None = None
+    font: str | ColumnExpr | None = None
+    size: str | ColumnExpr | None = None
+    align: Literal["center", "left", "right", "justify"] | ColumnExpr | None = None
+    v_align: Literal["middle", "top", "bottom"] | ColumnExpr | None = None
+    style: Literal["normal", "italic", "oblique"] | ColumnExpr | None = None
+    weight: Literal["normal", "bold", "bolder", "lighter"] | ColumnExpr | None = None
     stretch: Literal[
         "normal",
         "condensed",
@@ -92,12 +192,14 @@ class CellStyleText(CellStyle):
         "expanded",
         "extra-expanded",
         "ultra-expanded",
-    ] | None = None
-    decorate: Literal["overline", "line-through", "underline", "underline overline"] | None = None
-    transform: Literal["uppercase", "lowercase", "capitalize"] | None = None
+    ] | ColumnExpr | None = None
+    decorate: Literal[
+        "overline", "line-through", "underline", "underline overline"
+    ] | ColumnExpr | None = None
+    transform: Literal["uppercase", "lowercase", "capitalize"] | ColumnExpr | None = None
     whitespace: Literal[
         "normal", "nowrap", "pre", "pre-wrap", "pre-line", "break-spaces"
-    ] | None = None
+    ] | ColumnExpr | None = None
 
     def _to_html_style(self) -> str:
         rendered = ""
@@ -149,7 +251,7 @@ class CellStyleFill(CellStyle):
         value.
     """
 
-    color: str
+    color: str | ColumnExpr
     # alpha: Optional[float] = None
 
     def _to_html_style(self) -> str:
@@ -188,10 +290,10 @@ class CellStyleBorders(CellStyle):
 
     sides: Literal["all", "top", "bottom", "left", "right"] | List[
         Literal["all", "top", "bottom", "left", "right"]
-    ] = "all"
-    color: str = "#000000"
-    style: str = "solid"
-    weight: str = "1px"
+    ] | ColumnExpr = "all"
+    color: str | ColumnExpr = "#000000"
+    style: str | ColumnExpr = "solid"
+    weight: str | ColumnExpr = "1px"
 
     def _to_html_style(self) -> str:
         # If sides is an empty list, return an empty string

great_tables/_tbl_data.py

@@ -400,3 +400,46 @@ def _(ser: PdSeries) -> List[Any]:
 @to_list.register
 def _(ser: PlSeries) -> List[Any]:
     return ser.to_list()
+
+
+# mutate ----
+
+
+@singledispatch
+def eval_transform(df: DataFrameLike, expr: Any) -> List[Any]:
+    raise NotImplementedError(f"Unsupported type: {type(df)}")
+
+
+@eval_transform.register
+def _(df: PdDataFrame, expr: Callable[[PdDataFrame], PdSeries]) -> List[Any]:
+    res = expr(df)
+
+    if not isinstance(res, PdSeries):
+        raise ValueError(f"Result must be a pandas Series. Received {type(res)}")
+    elif not len(res) == len(df):
+        raise ValueError(
+            f"Result must be same length as input data. Observed different lengths."
+            f"\n\nInput data: {len(df)}.\nResult: {len(res)}."
+        )
+
+    return res.to_list()
+
+
+@eval_transform.register
+def _(df: PlDataFrame, expr: PlExpr) -> List[Any]:
+    df_res = df.select(expr)
+
+    if len(df_res.columns) > 1:
+        raise ValueError(f"Result must be a single column. Received {len(df_res.columns)} columns.")
+    else:
+        res = df_res[df_res.columns[0]]
+
+    if not isinstance(res, PlSeries):
+        raise ValueError(f"Result must be a polars Series. Received {type(res)}")
+    elif not len(res) == len(df):
+        raise ValueError(
+            f"Result must be same length as input data. Observed different lengths."
+            f"\n\nInput data: {len(df)}.\nResult: {len(res)}."
+        )
+
+    return res.to_list()

tests/__snapshots__/test_locations.ambr

@@ -0,0 +1,9 @@
+# serializer version: 1
+# name: test_set_style_loc_title_from_column_error
+  '''
+  Location type <class 'great_tables._locations.LocTitle'> cannot use FromColumn.
+
+  Style type: <class 'great_tables._styles.CellStyleText'>
+  Field with FromColumn: color
+  '''
+# ---

tests/test_locations.py

@@ -1,15 +1,19 @@
 import pandas as pd
+import polars as pl
 import pytest
 
 from great_tables._locations import (
     LocColumnSpanners,
     LocBody,
+    LocTitle,
     CellPos,
     resolve,
     resolve_vector_i,
     resolve_cols_i,
     resolve_rows_i,
+    set_style,
 )
+from great_tables._styles import CellStyleText, FromColumn
 from great_tables._gt_data import Spanners, SpannerInfo
 from great_tables import GT
 
@@ -84,3 +88,46 @@ def test_resolve_column_spanners_error_missing():
 
     with pytest.raises(ValueError):
         resolve(loc, spanners)
+
+
+@pytest.mark.parametrize(
+    "expr",
+    [
+        FromColumn("color"),
+        pl.col("color"),
+        pl.col("color").str.to_uppercase().str.to_lowercase(),
+    ],
+)
+def test_set_style_loc_body_from_column(expr):
+    df = pd.DataFrame({"x": [1, 2], "color": ["red", "blue"]})
+
+    if isinstance(expr, pl.Expr):
+        gt_df = GT(pl.DataFrame(df))
+    else:
+        gt_df = GT(df)
+
+    loc = LocBody(["x"], [1])
+    style = CellStyleText(color=expr)
+
+    new_gt = set_style(loc, gt_df, [style])
+
+    # 1 style info added
+    assert len(new_gt._styles) == 1
+    cell_info = new_gt._styles[0]
+
+    # style info has single cell style, with new color
+    assert len(cell_info.styles) == 1
+    assert isinstance(cell_info.styles[0], CellStyleText)
+    assert cell_info.styles[0].color == "blue"
+
+
+def test_set_style_loc_title_from_column_error(snapshot):
+    df = pd.DataFrame({"x": [1, 2], "color": ["red", "blue"]})
+    gt_df = GT(df)
+    loc = LocTitle("title")
+    style = CellStyleText(color=FromColumn("color"))
+
+    with pytest.raises(TypeError) as exc_info:
+        set_style(loc, gt_df, [style])
+
+    assert snapshot == exc_info.value.args[0]

tests/test_styles.py

@@ -0,0 +1,26 @@
+import pandas as pd
+
+from great_tables._styles import FromColumn, CellStyleText
+
+
+def test_from_column_replace():
+    """FromColumn is replaced by the specified column's value in a row of data"""
+
+    df = pd.DataFrame({"x": [1, 2], "color": ["red", "blue"]})
+    from_col = FromColumn("color")
+
+    style = CellStyleText(color=from_col)
+    new_style = style._from_row(df, 0)
+
+    assert style.color is from_col
+    assert new_style.color == "red"
+
+
+def test_from_column_fn():
+    df = pd.DataFrame({"x": [1, 2], "color": ["red", "blue"]})
+    from_col = FromColumn("color", fn=lambda x: x.upper())
+
+    style = CellStyleText(color=from_col)
+    new_style = style._from_row(df, 0)
+
+    assert new_style.color == "RED"