Commit 9000d416 authored by Floreal Cabanettes's avatar Floreal Cabanettes
Browse files

Add tooltree class

parent 83574b33
import random, string
class ToolTree:
@staticmethod
def randomword(length=10):
return ''.join(random.choice(string.ascii_lowercase) for i in range(length))
def __init__(self, name="", choices=list()):
self.name = name
self.id = self.randomword()
self.choices = choices if name == "" else []
self.childs = []
self.parents = []
def add_tool(self, tool_name: str):
new_tool = ToolTree(name=tool_name)
self.add_child(new_tool)
return new_tool
def add_choices(self, *choices: str):
choices = list(choices)
if len(choices) >= 2:
new_tool = ToolTree(choices=choices)
self.add_child(new_tool)
return new_tool
raise Exception("You must specify at least two choices in tools description")
def add_child(self, node: "ToolTree"):
self.childs.append(node)
node.parents.append(self)
def build_printable_tree(self, printed_parents=dict()):
nodes = []
child_nodes = []
for child in self.childs:
nb_parents = len(child.parents)
if nb_parents == 1 or (child.id in printed_parents and len(printed_parents[child.id]) == nb_parents):
if child.name != "":
nodes.append(child.name)
elif len(child.choices) > 0:
nodes.append(" OR ".join(child.choices))
for my_child in child.childs:
if my_child.id not in printed_parents:
printed_parents[my_child.id] = []
printed_parents[my_child.id].append(child.id)
new_child_nodes, printed_parents = child.build_printable_tree()
child_nodes += new_child_nodes
nodes += child_nodes
return nodes, printed_parents
def get_printable_tree(self):
nodes, printed_parents = self.build_printable_tree()
return nodes
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment