Merge pull request #1427 from JoshKarpel/clear-table-cols

Add option to clear columns in `DataTable.clear()`

Author: willmcgugan

CHANGELOG.md

@@ -5,6 +5,12 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](http://keepachangelog.com/)
 and this project adheres to [Semantic Versioning](http://semver.org/).
 
+## [0.8.3] - Unreleased
+
+### Added
+
+- Added an option to clear columns in DataTable.clear() https://github.com/Textualize/textual/pull/1427
+
 ## [0.8.2] - 2022-12-28
 
 ### Fixed

src/textual/widgets/_data_table.py

@@ -312,7 +312,7 @@ def _get_cell_region(self, row_index: int, column_index: int) -> Region:
         cell_region = Region(x, y, width, height)
         return cell_region
 
-    def clear(self) -> None:
+    def clear(self, columns: bool = False) -> None:
         """Clear the table.
 
         Args:
@@ -323,6 +323,8 @@ def clear(self) -> None:
         self._y_offsets.clear()
         self.data.clear()
         self.rows.clear()
+        if columns:
+            self.columns.clear()
         self._line_no = 0
         self._require_update_dimensions = True
         self.refresh()

tests/test_table.py

@@ -1,5 +1,7 @@
 import asyncio
 
+from rich.text import Text
+
 from textual.app import App, ComposeResult
 from textual.widgets import DataTable
 
@@ -18,13 +20,32 @@ async def test_table_clear() -> None:
         table.add_columns("foo", "bar")
         assert table.row_count == 0
         table.add_row("Hello", "World!")
+        assert [col.label for col in table.columns] == [Text("foo"), Text("bar")]
         assert table.data == {0: ["Hello", "World!"]}
         assert table.row_count == 1
         table.clear()
+        assert [col.label for col in table.columns] == [Text("foo"), Text("bar")]
         assert table.data == {}
         assert table.row_count == 0
 
 
+async def test_table_clear_with_columns() -> None:
+    """Check DataTable.clear(columns=True)"""
+
+    app = TableApp()
+    async with app.run_test() as pilot:
+        table = app.query_one(DataTable)
+        table.add_columns("foo", "bar")
+        assert table.row_count == 0
+        table.add_row("Hello", "World!")
+        assert [col.label for col in table.columns] == [Text("foo"), Text("bar")]
+        assert table.data == {0: ["Hello", "World!"]}
+        assert table.row_count == 1
+        table.clear(columns=True)
+        assert [col.label for col in table.columns] == []
+        assert table.data == {}
+        assert table.row_count == 0
+
 async def test_table_add_row() -> None:
 
     app = TableApp()