Witllm/wit/dataset/node_tree.py

from anytree import RenderTree

from anytree.node.nodemixin import NodeMixin
from anytree.node.util import _repr


class NodeTree(NodeMixin):
    def __init__(self, name, parent=None, children=None, **kwargs):
        self.__dict__.update(kwargs)
        self.name = name
        self.prop = ""
        self.parent = parent
        if children:
            self.children = children
        self.seq_node = []

    def __repr__(self):
        args = ["%r" % self.separator.join([""] + [str(node.name) for node in self.path])]
        return _repr(self, args=args, nameblacklist=["name"])

    def set_seq_prop(self, index, prop):
        self.seq_node[index].prop = prop

    def print(self):
        treestr = ""
        for pre, fill, node in RenderTree(self):
            treestr += f"{pre}{node.name} {node.prop}\n"
        print(treestr)