chore: use GTSelf for cols_ funcs, add cols_hide test

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

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)