feat(tree): enable add at location (#4772)

* enable add before or after index

* enable add before or after node

* try to improve error message wording

* raise typerror if invalid argument

* add new params to docstring

* add raises and note to docstring

* add before and after params to add_leaf

* update changelog

* fix copypasta in docstring

* improve error message wording

Co-authored by: Darren Burns <[email protected]>

---------

Co-authored-by: Will McGugan <[email protected]>

CHANGELOG.md

@@ -19,6 +19,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 - Add `Tree.move_cursor` to programmatically move the cursor without selecting the node https://github.com/Textualize/textual/pull/4753
 - Added `Footer` component style handling of padding for the key/description https://github.com/Textualize/textual/pull/4651
 - `StringKey` is now exported from `data_table` https://github.com/Textualize/textual/pull/4760
+- `TreeNode.add` and `TreeNode.add_leaf` now accepts `before` and `after` arguments to position a new node https://github.com/Textualize/textual/pull/4772
 - Added a `gradient` parameter to the `ProgressBar` widget https://github.com/Textualize/textual/pull/4774
 
 ### Fixed

src/textual/widgets/_tree.py

@@ -50,6 +50,10 @@ class UnknownNodeID(Exception):
     """Exception raised when referring to an unknown [`TreeNode`][textual.widgets.tree.TreeNode] ID."""
 
 
+class AddNodeError(Exception):
+    """Exception raised when there is an error with a request to add a node."""
+
+
 @dataclass
 class _TreeLine(Generic[TreeDataType]):
     path: list[TreeNode[TreeDataType]]
@@ -322,6 +326,8 @@ def add(
         label: TextType,
         data: TreeDataType | None = None,
         *,
+        before: int | TreeNode[TreeDataType] | None = None,
+        after: int | TreeNode[TreeDataType] | None = None,
         expand: bool = False,
         allow_expand: bool = True,
     ) -> TreeNode[TreeDataType]:
@@ -330,34 +336,101 @@ def add(
         Args:
             label: The new node's label.
             data: Data associated with the new node.
+            before: Optional index or `TreeNode` to add the node before.
+            after: Optional index or `TreeNode` to add the node after.
             expand: Node should be expanded.
             allow_expand: Allow use to expand the node via keyboard or mouse.
 
         Returns:
             A new Tree node
+
+        Raises:
+            AddNodeError: If there is a problem with the addition request.
+
+        Note:
+            Only one of `before` or `after` can be provided. If both are
+            provided a `AddNodeError` will be raised.
         """
+        if before is not None and after is not None:
+            raise AddNodeError("Unable to add a node both before and after a node")
+
+        insert_index: int = len(self.children)
+
+        if before is not None:
+            if isinstance(before, int):
+                insert_index = before
+            elif isinstance(before, TreeNode):
+                try:
+                    insert_index = self.children.index(before)
+                except ValueError:
+                    raise AddNodeError(
+                        "The node specified for `before` is not a child of this node"
+                    )
+            else:
+                raise TypeError(
+                    "`before` argument must be an index or a TreeNode object to add before"
+                )
+
+        if after is not None:
+            if isinstance(after, int):
+                insert_index = after + 1
+                if after < 0:
+                    insert_index += len(self.children)
+            elif isinstance(after, TreeNode):
+                try:
+                    insert_index = self.children.index(after) + 1
+                except ValueError:
+                    raise AddNodeError(
+                        "The node specified for `after` is not a child of this node"
+                    )
+            else:
+                raise TypeError(
+                    "`after` argument must be an index or a TreeNode object to add after"
+                )
+
         text_label = self._tree.process_label(label)
         node = self._tree._add_node(self, text_label, data)
         node._expanded = expand
         node._allow_expand = allow_expand
         self._updates += 1
-        self._children.append(node)
+        self._children.insert(insert_index, node)
         self._tree._invalidate()
         return node
 
     def add_leaf(
-        self, label: TextType, data: TreeDataType | None = None
+        self,
+        label: TextType,
+        data: TreeDataType | None = None,
+        *,
+        before: int | TreeNode[TreeDataType] | None = None,
+        after: int | TreeNode[TreeDataType] | None = None,
     ) -> TreeNode[TreeDataType]:
         """Add a 'leaf' node (a node that can not expand).
 
         Args:
             label: Label for the node.
             data: Optional data.
+            before: Optional index or `TreeNode` to add the node before.
+            after: Optional index or `TreeNode` to add the node after.
 
         Returns:
             New node.
+
+        Raises:
+            AddNodeError: If there is a problem with the addition request.
+
+        Note:
+            Only one of `before` or `after` can be provided. If both are
+            provided a `AddNodeError` will be raised.
         """
-        node = self.add(label, data, expand=False, allow_expand=False)
+        node = self.add(
+            label,
+            data,
+            before=before,
+            after=after,
+            expand=False,
+            allow_expand=False,
+        )
         return node
 
     def _remove_children(self) -> None:

src/textual/widgets/tree.py

@@ -1,6 +1,7 @@
 """Make non-widget Tree support classes available."""
 
 from ._tree import (
+    AddNodeError,
     EventTreeDataType,
     NodeID,
     RemoveRootError,
@@ -10,6 +11,7 @@
 )
 
 __all__ = [
+    "AddNodeError",
     "EventTreeDataType",
     "NodeID",
     "RemoveRootError",

tests/tree/test_tree_node_add.py

@@ -0,0 +1,120 @@
+import pytest
+
+from textual.widgets import Tree
+from textual.widgets.tree import AddNodeError
+
+
+def test_tree_node_add_before_and_after_raises_exception():
+    tree = Tree[None]("root")
+    with pytest.raises(AddNodeError):
+        tree.root.add("error", before=99, after=0)
+
+
+def test_tree_node_add_before_or_after_with_invalid_type_raises_exception():
+    tree = Tree[None]("root")
+    tree.root.add("node")
+    with pytest.raises(TypeError):
+        tree.root.add("before node", before="node")
+    with pytest.raises(TypeError):
+        tree.root.add("after node", after="node")
+
+
+def test_tree_node_add_before_index():
+    tree = Tree[None]("root")
+    tree.root.add("node")
+    tree.root.add("before node", before=0)
+    tree.root.add("first", before=-99)
+    tree.root.add("after first", before=-2)
+    tree.root.add("last", before=99)
+    tree.root.add("after node", before=4)
+    tree.root.add("before last", before=-1)
+
+    assert str(tree.root.children[0].label) == "first"
+    assert str(tree.root.children[1].label) == "after first"
+    assert str(tree.root.children[2].label) == "before node"
+    assert str(tree.root.children[3].label) == "node"
+    assert str(tree.root.children[4].label) == "after node"
+    assert str(tree.root.children[5].label) == "before last"
+    assert str(tree.root.children[6].label) == "last"
+
+
+def test_tree_node_add_after_index():
+    tree = Tree[None]("root")
+    tree.root.add("node")
+    tree.root.add("after node", after=0)
+    tree.root.add("first", after=-99)
+    tree.root.add("after first", after=-3)
+    tree.root.add("before node", after=1)
+    tree.root.add("before last", after=99)
+    tree.root.add("last", after=-1)
+
+    assert str(tree.root.children[0].label) == "first"
+    assert str(tree.root.children[1].label) == "after first"
+    assert str(tree.root.children[2].label) == "before node"
+    assert str(tree.root.children[3].label) == "node"
+    assert str(tree.root.children[4].label) == "after node"
+    assert str(tree.root.children[5].label) == "before last"
+    assert str(tree.root.children[6].label) == "last"
+
+
+def test_tree_node_add_relative_to_unknown_node_raises_exception():
+    tree = Tree[None]("root")
+    removed_node = tree.root.add("removed node")
+    removed_node.remove()
+    with pytest.raises(AddNodeError):
+        tree.root.add("node", before=removed_node)
+    with pytest.raises(AddNodeError):
+        tree.root.add("node", after=removed_node)
+
+
+def test_tree_node_add_before_node():
+    tree = Tree[None]("root")
+    node = tree.root.add("node")
+    before_node = tree.root.add("before node", before=node)
+    tree.root.add("first", before=before_node)
+    tree.root.add("after first", before=before_node)
+    last = tree.root.add("last", before=4)
+    before_last = tree.root.add("before last", before=last)
+    tree.root.add("after node", before=before_last)
+
+    assert str(tree.root.children[0].label) == "first"
+    assert str(tree.root.children[1].label) == "after first"
+    assert str(tree.root.children[2].label) == "before node"
+    assert str(tree.root.children[3].label) == "node"
+    assert str(tree.root.children[4].label) == "after node"
+    assert str(tree.root.children[5].label) == "before last"
+    assert str(tree.root.children[6].label) == "last"
+
+
+def test_tree_node_add_after_node():
+    tree = Tree[None]("root")
+    node = tree.root.add("node")
+    after_node = tree.root.add("after node", after=node)
+    first = tree.root.add("first", after=-3)
+    after_first = tree.root.add("after first", after=first)
+    tree.root.add("before node", after=after_first)
+    before_last = tree.root.add("before last", after=after_node)
+    tree.root.add("last", after=before_last)
+
+    assert str(tree.root.children[0].label) == "first"
+    assert str(tree.root.children[1].label) == "after first"
+    assert str(tree.root.children[2].label) == "before node"
+    assert str(tree.root.children[3].label) == "node"
+    assert str(tree.root.children[4].label) == "after node"
+    assert str(tree.root.children[5].label) == "before last"
+    assert str(tree.root.children[6].label) == "last"
+
+
+def test_tree_node_add_leaf_before_or_after():
+    tree = Tree[None]("root")
+    leaf = tree.root.add_leaf("leaf")
+    tree.root.add_leaf("before leaf", before=leaf)
+    tree.root.add_leaf("after leaf", after=leaf)
+    tree.root.add_leaf("first", before=0)
+    tree.root.add_leaf("last", after=-1)
+
+    assert str(tree.root.children[0].label) == "first"
+    assert str(tree.root.children[1].label) == "before leaf"
+    assert str(tree.root.children[2].label) == "leaf"
+    assert str(tree.root.children[3].label) == "after leaf"
+    assert str(tree.root.children[4].label) == "last"