Support mask= argument in LocBody

great_tables/_locations.py

@@ -539,6 +539,7 @@
 
     columns: SelectExpr = None
     rows: RowSelectExpr = None
+    mask: PlExpr | None = None
 
 
 @dataclass
@@ -823,6 +824,41 @@
     )
 
 
+def resolve_mask_i(
+    data: GTData | list[str],
+    expr: PlExpr,
+    excl_stub: bool = True,
+    excl_group: bool = True,
+) -> list[tuple[int, int, str]]:
+    """Return data for creating `CellPos`, based on expr"""
+    if not isinstance(expr, PlExpr):
+        raise ValueError("Only Polars expressions can be passed to the `mask` argument.")
+
+    stub_var = data._boxhead.vars_from_type(ColInfoTypeEnum.stub)
+    group_var = data._boxhead.vars_from_type(ColInfoTypeEnum.row_group)
+    cols_excl = [*(stub_var if excl_stub else []), *(group_var if excl_group else [])]
+
+    frame: PlDataFrame = data._tbl_data
+    df = frame.select(expr)
+
+    newer_row_attr, older_row_attr = "with_row_index", "with_row_count"
+    row_number_colname = "__row_number__"
+    row_attr = newer_row_attr if hasattr(df, newer_row_attr) else older_row_attr
+
+    # Add row numbers after `df.select()`, as the `__row_number__` column type is `UInt32`.
+    df_with_row_number = getattr(df, row_attr)(name=row_number_colname)
+
+    select_columns = [col for col in df.columns if col not in cols_excl]
+
+    cellpos_data: list[tuple[int, int, str]] = []  # column, row, colname for `CellPos`
+    for row_dict in df_with_row_number.iter_rows(named=True):
+        for col_idx, colname in enumerate(select_columns):
+            if row_dict[colname]:  # select only when `row_dict[colname]` is True
+                row_idx = row_dict[row_number_colname]
+                cellpos_data.append((col_idx, row_idx, colname))
+    return cellpos_data
+
+
 # Resolve generic ======================================================================
 
 
@@ -869,14 +905,21 @@
 @resolve.register
 def _(loc: LocBody, data: GTData) -> list[CellPos]:
     cols = resolve_cols_i(data=data, expr=loc.columns)
-    rows = resolve_rows_i(data=data, expr=loc.rows)
-
-    # TODO: dplyr arranges by `Var1`, and does distinct (since you can tidyselect the same
-    # thing multiple times
-    cell_pos = [
-        CellPos(col[1], row[1], colname=col[0]) for col, row in itertools.product(cols, rows)
-    ]
+    if (loc.columns is not None or loc.rows is not None) and loc.mask is not None:
+        raise ValueError(
+            "Cannot specify the `mask` argument along with `columns` or `rows` in `loc.body()`."
+        )
 
+    if loc.mask is None:
+        rows = resolve_rows_i(data=data, expr=loc.rows)
+        # TODO: dplyr arranges by `Var1`, and does distinct (since you can tidyselect the same
+        # thing multiple times
+        cell_pos = [
+            CellPos(col[1], row[1], colname=col[0]) for col, row in itertools.product(cols, rows)
+        ]
+    else:
+        cellpos_data = resolve_mask_i(data=data, expr=loc.mask)
+        cell_pos = [CellPos(*cellpos) for cellpos in cellpos_data]
     return cell_pos
 
 

tests/test_tab_create_modify.py

@@ -5,6 +5,7 @@
 from great_tables._locations import LocBody
 from great_tables._styles import CellStyleFill
 from great_tables._tab_create_modify import tab_style
+from polars import selectors as cs
 
 
 @pytest.fixture
@@ -33,6 +34,30 @@ def test_tab_style_multiple_columns(gt: GT):
     assert new_gt._styles[0].styles[0] is style
 
 
+def test_tab_style_loc_body_mask():
+    gt = GT(pl.DataFrame({"x": [1, 2], "y": [4, 5]}))
+    style = CellStyleFill(color="blue")
+    new_gt = tab_style(gt, style, LocBody(mask=cs.numeric().gt(1.5)))
+
+    assert len(gt._styles) == 0
+    assert len(new_gt._styles) == 3
+
+    xy_0y, xy_1x, xy_1y = new_gt._styles
+
+    assert xy_0y.styles[0] is style
+    assert xy_1x.styles[0] is style
+    assert xy_1y.styles[0] is style
+
+    assert xy_0y.rownum == 0
+    assert xy_0y.colname == "y"
+
+    assert xy_1x.rownum == 1
+    assert xy_1x.colname == "x"
+
+    assert xy_1y.rownum == 1
+    assert xy_1y.colname == "y"
+
+
 def test_tab_style_google_font(gt: GT):
     new_gt = tab_style(
         gt,