Modify the locations parameter of opt_all_caps() to accept Loc objects

Author: jrycw

great_tables/_options.py

@@ -893,8 +893,8 @@ def opt_all_caps(
 
     locations
         Which locations should undergo this text transformation? By default it includes all of
-        the `"column_labels"`, the `"stub"`, and the `"row_group"` locations. However, we could
-        just choose one or two of those.
+        the `loc.column_labels`, the `loc.stub"`, and the `loc.row_group` locations. However, we
+        could just choose one or two of those.
 
     Returns
     -------
@@ -909,7 +909,7 @@ def opt_all_caps(
     in all row groups is transformed to all caps using the `opt_all_caps()` method.
 
     ```{python}
-    from great_tables import GT, exibble, md
+    from great_tables import GT, exibble, loc, md
 
     (
       GT(
@@ -927,16 +927,46 @@ def opt_all_caps(
       .opt_all_caps()
     )
     ```
+    `opt_all_caps()` accepts a `locations` parameter that allows us to specify which components
+    should be transformed. For example, if we only want to ensure that all text in the stub and all
+    row groups is converted to all caps:
+    ```{python}
+    (
+      GT(
+        exibble[["num", "char", "currency", "row", "group"]],
+        rowname_col="row",
+        groupname_col="group"
+      )
+      .tab_header(
+        title=md("Data listing from **exibble**"),
+        subtitle=md("`exibble` is a **Great Tables** dataset.")
+      )
+      .fmt_number(columns="num")
+      .fmt_currency(columns="currency")
+      .tab_source_note(source_note="This is only a subset of the dataset.")
+      .opt_all_caps(locations=[loc.stub, loc.row_group])
+    )
+    ```
     """
+    from great_tables._locations import Loc, LocColumnLabels, LocStub, LocRowGroups
 
-    # If providing a scalar string value, normalize it to be in a list
-    if not isinstance(locations, list):
-        locations = _utils._str_scalar_to_list(cast(str, locations))
-
-    # Ensure that the `locations` value is a list of strings
-    _utils._assert_str_list(locations)
+    if not locations:
+        locations = [LocColumnLabels, LocStub, LocRowGroups]
 
-    # TODO: Ensure that all values within `locations` are valid
+    # If providing a Loc object, normalize it to be in a list
+    if not isinstance(locations, list):
+        locations = [locations]
+
+    # Ensure that all values within `locations` are valid
+    # A `try-except` block is needed here because the first argument of `issubclass()` must be a
+    # class.
+    for location in locations:
+        try:
+            issubclass(location, Loc)
+        except TypeError:
+            raise AssertionError(
+                "Only `loc.column_labels`, `loc.stub` and `loc.row_group` are allowed in the locations."
+            )
 
     # if `all_caps` is False, reset options to default, or, set new options
     # for `locations` selected
@@ -956,23 +986,23 @@ def opt_all_caps(
 
     info = [
         (
-            "column_labels",
+            LocColumnLabels,
             {
                 "column_labels_font_size": "80%",
                 "column_labels_font_weight": "bolder",
                 "column_labels_text_transform": "uppercase",
             },
         ),
         (
-            "stub",
+            LocStub,
             {
                 "stub_font_size": "80%",
                 "stub_font_weight": "bolder",
                 "stub_text_transform": "uppercase",
             },
         ),
         (
-            "row_group",
+            LocRowGroups,
             {
                 "row_group_font_size": "80%",
                 "row_group_font_weight": "bolder",

great_tables/loc.py

@@ -4,6 +4,7 @@
     LocBody as body,
     LocStub as stub,
     LocColumnLabels as column_labels,
+    LocRowGroups as row_group,
 )
 
-__all__ = ("body", "stub", "column_labels")
+__all__ = ("body", "stub", "column_labels", "row_group")

tests/test_options.py

@@ -1,6 +1,6 @@
 import pandas as pd
 import pytest
-from great_tables import GT, exibble, md
+from great_tables import GT, exibble, loc, md
 from great_tables._scss import compile_scss
 from great_tables._gt_data import default_fonts_list
 
@@ -328,7 +328,6 @@ def test_scss_from_opt_table_outline(gt_tbl: GT, snapshot):
 
 
 def test_opt_table_font_add_font():
-
     gt_tbl = GT(exibble).opt_table_font(font="Arial", weight="bold", style="italic")
 
     assert gt_tbl._options.table_font_names.value == ["Arial"] + default_fonts_list
@@ -369,3 +368,73 @@ def test_opt_table_font_raises():
         GT(exibble).opt_table_font(font=None, stack=None)
 
     assert "Either `font=` or `stack=` must be provided." in exc_info.value.args[0]
+
+
+def test_opt_all_caps(gt_tbl: GT):
+    tbl = gt_tbl.opt_all_caps(locations=loc.column_labels)
+
+    assert tbl._options.column_labels_font_size.value == "80%"
+    assert tbl._options.column_labels_font_weight.value == "bolder"
+    assert tbl._options.column_labels_text_transform.value == "uppercase"
+
+    tbl = gt_tbl.opt_all_caps(locations=[loc.column_labels, loc.stub])
+
+    assert tbl._options.column_labels_font_size.value == "80%"
+    assert tbl._options.column_labels_font_weight.value == "bolder"
+    assert tbl._options.column_labels_text_transform.value == "uppercase"
+
+    assert tbl._options.stub_font_size.value == "80%"
+    assert tbl._options.stub_font_weight.value == "bolder"
+    assert tbl._options.stub_text_transform.value == "uppercase"
+
+    tbl = gt_tbl.opt_all_caps(locations=[loc.column_labels, loc.stub, loc.row_group])
+
+    assert tbl._options.column_labels_font_size.value == "80%"
+    assert tbl._options.column_labels_font_weight.value == "bolder"
+    assert tbl._options.column_labels_text_transform.value == "uppercase"
+
+    assert tbl._options.stub_font_size.value == "80%"
+    assert tbl._options.stub_font_weight.value == "bolder"
+    assert tbl._options.stub_text_transform.value == "uppercase"
+
+    assert tbl._options.row_group_font_size.value == "80%"
+    assert tbl._options.row_group_font_weight.value == "bolder"
+    assert tbl._options.row_group_text_transform.value == "uppercase"
+
+    tbl = gt_tbl.opt_all_caps()
+
+    assert tbl._options.column_labels_font_size.value == "80%"
+    assert tbl._options.column_labels_font_weight.value == "bolder"
+    assert tbl._options.column_labels_text_transform.value == "uppercase"
+
+    assert tbl._options.stub_font_size.value == "80%"
+    assert tbl._options.stub_font_weight.value == "bolder"
+    assert tbl._options.stub_text_transform.value == "uppercase"
+
+    assert tbl._options.row_group_font_size.value == "80%"
+    assert tbl._options.row_group_font_weight.value == "bolder"
+    assert tbl._options.row_group_text_transform.value == "uppercase"
+
+    tbl = gt_tbl.opt_all_caps(all_caps=False)
+
+    assert tbl._options.column_labels_font_size.value == "100%"
+    assert tbl._options.column_labels_font_weight.value == "normal"
+    assert tbl._options.column_labels_text_transform.value == "inherit"
+
+    assert tbl._options.stub_font_size.value == "100%"
+    assert tbl._options.stub_font_weight.value == "initial"
+    assert tbl._options.stub_text_transform.value == "inherit"
+
+    assert tbl._options.row_group_font_size.value == "100%"
+    assert tbl._options.row_group_font_weight.value == "initial"
+    assert tbl._options.row_group_text_transform.value == "inherit"
+
+
+def test_opt_all_caps_raises(gt_tbl: GT):
+    with pytest.raises(AssertionError) as exc_info:
+        gt_tbl.opt_all_caps(locations="column_labels")
+
+    assert (
+        f"Only `loc.column_labels`, `loc.stub` and `loc.row_group` are allowed in the locations."
+        in exc_info.value.args[0]
+    )