Merge pull request #1644 from davep/tree-deeply

Add support for a method of expanding/collapsing all tree nodes from a given node down

CHANGELOG.md

@@ -9,6 +9,9 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 
 ### Added
 
+- Added `TreeNode.expand_all` https://github.com/Textualize/textual/issues/1430
+- Added `TreeNode.collapse_all` https://github.com/Textualize/textual/issues/1430
+- Added `TreeNode.toggle_all` https://github.com/Textualize/textual/issues/1430
 - Added the coroutines `Animator.wait_until_complete` and `pilot.wait_for_scheduled_animations` that allow waiting for all current and scheduled animations https://github.com/Textualize/textual/issues/1658
 - Added the method `Animator.is_being_animated` that checks if an attribute of an object is being animated or is scheduled for animation
 - Added more keyboard actions and related bindings to `Input` https://github.com/Textualize/textual/pull/1676

src/textual/widgets/_tree.py

@@ -159,23 +159,63 @@ class TreeNode(Generic[TreeDataType]):
         self._allow_expand = allow_expand
         self._updates += 1
 
-    def expand(self) -> None:
+    def _expand(self, expand_all: bool) -> None:
-        """Expand a node (show its children)."""
+        """Mark the node as expanded (its children are shown).
+
+        Args:
+            expand_all: If `True` expand all offspring at all depths.
+        """
         self._expanded = True
         self._updates += 1
+        if expand_all:
+            for child in self.children:
+                child._expand(expand_all)
+
+    def expand(self) -> None:
+        """Expand the node (show its children)."""
+        self._expand(False)
         self._tree._invalidate()
 
-    def collapse(self) -> None:
+    def expand_all(self) -> None:
-        """Collapse the node (hide children)."""
+        """Expand the node (show its children) and all those below it."""
+        self._expand(True)
+        self._tree._invalidate()
+
+    def _collapse(self, collapse_all: bool) -> None:
+        """Mark the node as collapsed (its children are hidden).
+
+        Args:
+            collapse_all: If `True` collapse all offspring at all depths.
+        """
         self._expanded = False
         self._updates += 1
+        if collapse_all:
+            for child in self.children:
+                child._collapse(collapse_all)
+
+    def collapse(self) -> None:
+        """Collapse the node (hide its children)."""
+        self._collapse(False)
+        self._tree._invalidate()
+
+    def collapse_all(self) -> None:
+        """Collapse the node (hide its children) and all those below it."""
+        self._collapse(True)
         self._tree._invalidate()
 
     def toggle(self) -> None:
-        """Toggle the expanded state."""
+        """Toggle the node's expanded state."""
-        self._expanded = not self._expanded
+        if self._expanded:
-        self._updates += 1
+            self.collapse()
-        self._tree._invalidate()
+        else:
+            self.expand()
+
+    def toggle_all(self) -> None:
+        """Toggle the node's expanded state and make all those below it match."""
+        if self._expanded:
+            self.collapse_all()
+        else:
+            self.expand_all()
 
     @property
     def label(self) -> TextType:

tests/tree/test_tree_expand_etc.py

@@ -0,0 +1,106 @@
+from __future__ import annotations
+
+from textual.app import App, ComposeResult
+from textual.widgets import Tree
+
+
+class TreeApp(App[None]):
+    """Test tree app."""
+
+    def compose(self) -> ComposeResult:
+        yield Tree("Test")
+
+    def on_mount(self) -> None:
+        tree = self.query_one(Tree)
+        for n in range(10):
+            tree.root.add(f"Trunk {n}")
+        node = tree.root.children[0]
+        for n in range(10):
+            node = node.add(str(n))
+
+
+async def test_tree_node_expand() -> None:
+    """Expanding one node should not expand all nodes."""
+    async with TreeApp().run_test() as pilot:
+        pilot.app.query_one(Tree).root.expand()
+        assert pilot.app.query_one(Tree).root.is_expanded is True
+        check_node = pilot.app.query_one(Tree).root.children[0]
+        while check_node.children:
+            assert any(child.is_expanded for child in check_node.children) is False
+            check_node = check_node.children[0]
+
+
+async def test_tree_node_expand_all() -> None:
+    """Expanding all on a node should expand all child nodes too."""
+    async with TreeApp().run_test() as pilot:
+        pilot.app.query_one(Tree).root.expand_all()
+        assert pilot.app.query_one(Tree).root.is_expanded is True
+        check_node = pilot.app.query_one(Tree).root.children[0]
+        while check_node.children:
+            assert check_node.children[0].is_expanded is True
+            assert any(child.is_expanded for child in check_node.children[1:]) is False
+            check_node = check_node.children[0]
+
+
+async def test_tree_node_collapse() -> None:
+    """Collapsing one node should not collapse all nodes."""
+    async with TreeApp().run_test() as pilot:
+        pilot.app.query_one(Tree).root.expand_all()
+        pilot.app.query_one(Tree).root.children[0].collapse()
+        assert pilot.app.query_one(Tree).root.children[0].is_expanded is False
+        check_node = pilot.app.query_one(Tree).root.children[0].children[0]
+        while check_node.children:
+            assert all(child.is_expanded for child in check_node.children) is True
+            check_node = check_node.children[0]
+
+
+async def test_tree_node_collapse_all() -> None:
+    """Collapsing all on a node should collapse all child noes too."""
+    async with TreeApp().run_test() as pilot:
+        pilot.app.query_one(Tree).root.expand_all()
+        pilot.app.query_one(Tree).root.children[0].collapse_all()
+        assert pilot.app.query_one(Tree).root.children[0].is_expanded is False
+        check_node = pilot.app.query_one(Tree).root.children[0].children[0]
+        while check_node.children:
+            assert check_node.children[0].is_expanded is False
+            assert all(child.is_expanded for child in check_node.children[1:]) is True
+            check_node = check_node.children[0]
+
+
+async def test_tree_node_toggle() -> None:
+    """Toggling one node should not toggle all nodes."""
+    async with TreeApp().run_test() as pilot:
+        assert pilot.app.query_one(Tree).root.is_expanded is False
+        check_node = pilot.app.query_one(Tree).root.children[0]
+        while check_node.children:
+            assert any(child.is_expanded for child in check_node.children) is False
+            check_node = check_node.children[0]
+        pilot.app.query_one(Tree).root.toggle()
+        assert pilot.app.query_one(Tree).root.is_expanded is True
+        check_node = pilot.app.query_one(Tree).root.children[0]
+        while check_node.children:
+            assert any(child.is_expanded for child in check_node.children) is False
+            check_node = check_node.children[0]
+
+
+async def test_tree_node_toggle_all() -> None:
+    """Toggling all on a node should toggle all child nodes too."""
+    async with TreeApp().run_test() as pilot:
+        assert pilot.app.query_one(Tree).root.is_expanded is False
+        check_node = pilot.app.query_one(Tree).root.children[0]
+        while check_node.children:
+            assert any(child.is_expanded for child in check_node.children) is False
+            check_node = check_node.children[0]
+        pilot.app.query_one(Tree).root.toggle_all()
+        assert pilot.app.query_one(Tree).root.is_expanded is True
+        check_node = pilot.app.query_one(Tree).root.children[0]
+        while check_node.children:
+            assert check_node.children[0].is_expanded is True
+            assert any(child.is_expanded for child in check_node.children[1:]) is False
+            check_node = check_node.children[0]
+        pilot.app.query_one(Tree).root.toggle_all()
+        assert pilot.app.query_one(Tree).root.is_expanded is False
+        check_node = pilot.app.query_one(Tree).root.children[0]
+        while check_node.children:
+            assert any(child.is_expanded for child in check_node.children) is False
+            check_node = check_node.children[0]