Merge pull request #86 from posit-dev/fix-cols-hide

fix: cols_hide

great_tables/_gt_data.py

@@ -267,12 +267,17 @@ def set_groupname_col(self, groupname_col: str):
                 new_col = replace(col, type=ColInfoTypeEnum.default)
                 self._d[ii] = new_col
 
-    def set_col_hidden(self, colname: str):
+    def set_cols_hidden(self, colnames: list[str]):
         # TODO: validate that colname is in the boxhead
+        res: list[ColInfo] = []
         for ii, col in enumerate(self._d):
-            if col.var == colname:
+            if col.var in colnames:
                 new_col = replace(col, type=ColInfoTypeEnum.hidden)
-                self._d[ii] = new_col
+                res.append(new_col)
+            else:
+                res.append(col)
+
+        return self.__class__(res)
 
     def align_from_data(self, data: TblData):
         """Updates align attribute in entries based on data types."""

great_tables/_spanners.py

@@ -8,7 +8,7 @@
 from ._locations import SelectExpr, resolve_cols_c
 
 if TYPE_CHECKING:
-    from ._gt_data import GTData, Boxhead
+    from ._gt_data import Boxhead
     from ._types import GTSelf
 
 
@@ -188,7 +188,7 @@ def tab_spanner(
     return new_data
 
 
-def cols_move(data: GTData, columns: SelectExpr, after: str) -> GTData:
+def cols_move(data: GTSelf, columns: SelectExpr, after: str) -> GTSelf:
     """Move one or more columns.
 
     On those occasions where you need to move columns this way or that way, we can make use of the
@@ -411,7 +411,7 @@ def cols_move_to_end(data: GTSelf, columns: SelectExpr) -> GTSelf:
     return data._replace(_boxhead=new_boxhead)
 
 
-def cols_hide(data: GTData, columns: SelectExpr) -> GTData:
+def cols_hide(data: GTSelf, columns: SelectExpr) -> GTSelf:
     """Hide one or more columns.
 
     The `cols_hide()` method allows us to hide one or more columns from appearing in the final
@@ -456,10 +456,7 @@ def cols_hide(data: GTData, columns: SelectExpr) -> GTData:
         raise ValueError("All `columns` must exist and be visible in the input `data` table.")
 
     # New boxhead with hidden columns
-    new_boxhead = data._boxhead
-
-    for col in sel_cols:
-        new_boxhead = new_boxhead.set_col_hidden(colname=col)
+    new_boxhead = data._boxhead.set_cols_hidden(sel_cols)
 
     return data._replace(_boxhead=new_boxhead)
 

tests/test_spanners.py

@@ -5,6 +5,7 @@
     spanners_print_matrix,
     empty_spanner_matrix,
     tab_spanner,
+    cols_hide,
     cols_move,
 )
 from great_tables._gt_data import Spanners, SpannerInfo, Boxhead, ColInfo, ColInfoTypeEnum
@@ -150,6 +151,17 @@ def test_tab_spanners_with_gather():
     assert [col.var for col in new_gt._boxhead] == ["a", "c", "b"]
 
 
+def test_cols_hide():
+    df = pd.DataFrame({"a": [1, 2], "b": [3, 4], "c": [5, 6]})
+    src_gt = GT(df)
+
+    new_gt = cols_hide(src_gt, columns=["a"])
+    assert [col.var for col in new_gt._boxhead if col.visible] == ["b", "c"]
+
+    new_gt = cols_hide(src_gt, columns=["a", "b"])
+    assert [col.var for col in new_gt._boxhead if col.visible] == ["c"]
+
+
 def test_cols_move():
     df = pd.DataFrame({"a": [1, 2], "b": [3, 4], "c": [5, 6]})
     src_gt = GT(df)