Witllm/wit/dataset/node_tree.py

29 lines
842 B
Python

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)