From 22d0d422517a7779a7d27bead223bf2701878fc5 Mon Sep 17 00:00:00 2001 From: TomJGooding <101601846+TomJGooding@users.noreply.github.com> Date: Thu, 18 Jul 2024 16:35:06 +0100 Subject: [PATCH] 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 --------- Co-authored-by: Will McGugan --- CHANGELOG.md | 1 + src/textual/widgets/_tree.py | 79 +++++++++++++++++++- src/textual/widgets/tree.py | 2 + tests/tree/test_tree_node_add.py | 120 +++++++++++++++++++++++++++++++ 4 files changed, 199 insertions(+), 3 deletions(-) create mode 100644 tests/tree/test_tree_node_add.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 25f79406ab..a70e754fdc 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/src/textual/widgets/_tree.py b/src/textual/widgets/_tree.py index ed44340d76..c424d9889d 100644 --- a/src/textual/widgets/_tree.py +++ b/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: diff --git a/src/textual/widgets/tree.py b/src/textual/widgets/tree.py index 70296bcaa2..d019168ab3 100644 --- a/src/textual/widgets/tree.py +++ b/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", diff --git a/tests/tree/test_tree_node_add.py b/tests/tree/test_tree_node_add.py new file mode 100644 index 0000000000..056ed58cbf --- /dev/null +++ b/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"