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(
+                f"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")