Skip to content

Commit

Permalink
initial add for the plugin
Browse files Browse the repository at this point in the history
remove fstring

migrate to Tree classmethod
  • Loading branch information
gaoyukid authored and 陈夏明 committed Aug 18, 2023
1 parent 8cd2940 commit 4162ce7
Show file tree
Hide file tree
Showing 2 changed files with 68 additions and 0 deletions.
34 changes: 34 additions & 0 deletions tests/test_tree.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,14 @@ def setUp(self):
# |-- George
self.tree = tree
self.copytree = Tree(self.tree, deep=True)
self.input_dict = {
"Bill": "Harry",
"Jane": "Harry",
"Harry": None,
"Diane": "Jane",
"Mark": "Jane",
"Mary": "Harry",
}

@staticmethod
def get_t1():
Expand Down Expand Up @@ -734,3 +742,29 @@ def test_root_removal(self):
t.create_node(identifier="root-B")
self.assertEqual(len(t.nodes.keys()), 1)
self.assertEqual(t.root, "root-B")

def test_from_map(self):
tree = Tree.from_map(self.input_dict)
self.assertTrue(tree.size() == 6)
self.assertTrue(
tree.root == [k for k, v in self.input_dict.items() if v is None][0]
)
tree = Tree.from_map(self.input_dict, id_func=lambda x: x.upper())
self.assertTrue(tree.size() == 6)

def data_func(x):
return x.upper()

tree = Tree.from_map(self.input_dict, data_func=data_func)
self.assertTrue(tree.size() == 6)
self.assertTrue(
tree.get_node(tree.root).data
== data_func([k for k, v in self.input_dict.items() if v is None][0])
)
with self.assertRaises(ValueError):
# invalid input payload without a root
tree = Tree.from_map({"a": "b"})

with self.assertRaises(ValueError):
# invalid input payload without more than 1 root
tree = Tree.from_map({"a": None, "b": None})
34 changes: 34 additions & 0 deletions treelib/tree.py
Original file line number Diff line number Diff line change
Expand Up @@ -1128,3 +1128,37 @@ def to_graphviz(
print(f.getvalue())

f.close()

@classmethod
def from_map(cls, child_parent_dict, id_func=None, data_func=None):
'''
takes a dict with child:parent, and form a tree
'''
tree = Tree()
if tree is None or tree.size() > 0:
raise ValueError('need to pass in an empty tree')
id_func = id_func if id_func else lambda x: x
data_func = data_func if data_func else lambda x: None
parent_child_dict = {}
root_node = None
for k, v in child_parent_dict.items():
if v is None and root_node is None:
root_node = k
elif v is None and root_node is not None:
raise ValueError('invalid input, more than 1 child has no parent')
else:
if v in parent_child_dict:
parent_child_dict[v].append(k)
else:
parent_child_dict[v] = [k]
if root_node is None:
raise ValueError('cannot find root')

tree.create_node(root_node, id_func(root_node), data=data_func(root_node))
queue = [root_node]
while len(queue) > 0:
parent_node = queue.pop()
for child in parent_child_dict.get(parent_node, []):
tree.create_node(child, id_func(child), parent=id_func(parent_node), data=data_func(child))
queue.append(child)
return tree

0 comments on commit 4162ce7

Please sign in to comment.