From 78f49f0ab2d87ec58589bc0e1082885f6e7c05db Mon Sep 17 00:00:00 2001
From: Dave Pearson <davep@davep.org>
Date: Wed, 1 Feb 2023 13:44:53 +0000
Subject: [PATCH] Allow replacement of the root node data when clearing the
 tree

---
 src/textual/widgets/_tree.py     |  6 ++++--
 tests/tree/test_tree_clearing.py | 24 +++++++++++++++++++++++-
 2 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/src/textual/widgets/_tree.py b/src/textual/widgets/_tree.py
index 8f2304ae8..f299c0d99 100644
--- a/src/textual/widgets/_tree.py
+++ b/src/textual/widgets/_tree.py
@@ -560,18 +560,20 @@ class Tree(Generic[TreeDataType], ScrollView, can_focus=True):
         label = self.render_label(node, NULL_STYLE, NULL_STYLE)
         return label.cell_len
 
-    def clear(self, label: TextType | None = None) -> None:
+    def clear(self, label: TextType | None = None, data: TreeDataType = ...) -> None:
         """Clear all nodes under root.
 
         Args:
             label: An optional new label for the root node. If not provided
                 the current root node's label will be used.
+            data: Optional new data for the root node. If not provided the
+                current root node's data will be used.
         """
         self._line_cache.clear()
         self._tree_lines_cached = None
         self._current_id = 0
         root_label = self.root._label if label is None else label
-        root_data = self.root.data
+        root_data = self.root.data if data is ... else data
         self.root = TreeNode(
             self,
             None,
diff --git a/tests/tree/test_tree_clearing.py b/tests/tree/test_tree_clearing.py
index 1a1efa449..9b7762bbf 100644
--- a/tests/tree/test_tree_clearing.py
+++ b/tests/tree/test_tree_clearing.py
@@ -56,7 +56,29 @@ async def test_tree_new_label_clear() -> None:
     async with TreeClearApp().run_test() as pilot:
         tree = pilot.app.query_one(VerseTree)
         assert len(tree.root.children) > 1
-        pilot.app.query_one(VerseTree).clear("Jiangyin")
+        pilot.app.query_one(VerseTree).clear(label="Jiangyin")
         assert len(tree.root.children) == 0
         assert str(tree.root.label) == "Jiangyin"
         assert isinstance(tree.root.data, VerseStar)
+
+
+async def test_tree_new_data_clear() -> None:
+    """Clearing a tree with data should keep the old label and use the new data."""
+    async with TreeClearApp().run_test() as pilot:
+        tree = pilot.app.query_one(VerseTree)
+        assert len(tree.root.children) > 1
+        pilot.app.query_one(VerseTree).clear(data=VersePlanet())
+        assert len(tree.root.children) == 0
+        assert str(tree.root.label) == "White Sun"
+        assert isinstance(tree.root.data, VersePlanet)
+
+
+async def test_tree_new_labal_and_data_clear() -> None:
+    """Clearing a tree with label and data should replace the label and data."""
+    async with TreeClearApp().run_test() as pilot:
+        tree = pilot.app.query_one(VerseTree)
+        assert len(tree.root.children) > 1
+        pilot.app.query_one(VerseTree).clear(label="Jiangyin", data=VersePlanet())
+        assert len(tree.root.children) == 0
+        assert str(tree.root.label) == "Jiangyin"
+        assert isinstance(tree.root.data, VersePlanet)