Adding new nodes by id in anytree

1.5k Views Asked by At

I'm trying to implement a tree with unique IDs per 'section' nodes and unique IDs for each child node in each section. I wonder how to accomplish this with anytree (python). I tried this:

from anytree import AnyNode, RenderTree, Resolver


class MyTree:

    def create(self):
        self.root = AnyNode(id="root")
        initial_section_node = AnyNode(id="section1", parent=self.root)
        AnyNode(id="sub0A", parent=initial_section_node)
        AnyNode(id="sub0B", parent=initial_section_node)
        another_section_node = AnyNode(id="section2", parent=self.root)
        AnyNode(id="sub0A", parent=another_section_node)
        AnyNode(id="sub0B", parent=another_section_node)

    def add_node_to_parent(self, section_id, node_id: str):
        r = Resolver("id")
        section_node = r.get(self.root, section_id)
        if node_id not in section_node.children:
            AnyNode(id=node_id, parent=section_node)
        else:
            print(node_id + " already exists")

    def display(self):
        print(RenderTree(self.root))


test_tree = MyTree()
test_tree.create()
test_tree.add_node_to_parent("section1", "sub0C")
test_tree.add_node_to_parent("section1", "sub0A")  # this will add another node with this id, but I want to prevent this
test_tree.display()

but it displays

AnyNode(id='root')
├── AnyNode(id='section1')
│   ├── AnyNode(id='sub0A')
│   ├── AnyNode(id='sub0B')
│   ├── AnyNode(id='sub0C')
│   └── AnyNode(id='sub0A')
└── AnyNode(id='section2')
    ├── AnyNode(id='sub0A')
    └── AnyNode(id='sub0B')

i.e., I can use the resolver to find the right section and then add new nodes to it, but the condition if node_id not in section_node.children: does not work because node_id is not a node object but a string object.

I understand that I need using the resolver once again to check if there is a child node of the "section1" node with the ID I want to add, but I would prefer using the collection section_node.children instead, because I think it works faster than the resolver, especially for big trees.

2

There are 2 best solutions below

1
On

As you said, you cannot directly check for inclusion of a string in a tuple of objects. You may define a function, e.g.

if not list(filter(lambda x:x.id==node.id, section_node.children)):

But I doubt this can be faster than Resolver.

Alternatively you may handle a custom attribute children_ids, and use that for your check.

1
On

I found a possible solution. We can use the found section node in the resolver (instead of the whole root node) to look for child nodes. The modified add_node_to_parent function looks like this:

def add_node_to_parent(self, section_id, node_id: str):
    r = Resolver("id")
    section_node = r.get(self.root, section_id)
    try:
        node = r.get(section_node, node_id)
        print(node_id + " already exists")
    except anytree.resolver.ChildResolverError:
        AnyNode(id=node_id, parent=section_node)

Now, the output is as expected:

sub0A already exists
AnyNode(id='root')
├── AnyNode(id='section1')
│   ├── AnyNode(id='sub0A')
│   ├── AnyNode(id='sub0B')
│   └── AnyNode(id='sub0C')
└── AnyNode(id='section2')
    ├── AnyNode(id='sub0A')
    └── AnyNode(id='sub0B')

Of course, the line r.get(self.root, section_id) can also raise a ChildResolverError exception if section_id is not found but the answer demonstrates the principle.