def test_reduce_to_one_leaf(self):
# reduce to A
leaves = {'A'}
expected_tree = loads("{A};")
cutter = TreeCutter()
cutter.attach_tree(self.tree)
cutter.choose_leaves(leaves)
result = cutter.cut()
compare_nodes(self, result.get_newick()[0], expected_tree[0])
def test_reduce_to_sub_branch(self):
# reduce to {E F}
leaves = {'E', 'F'}
expected_tree = loads("({E},{F}){E F};")
cutter = TreeCutter()
cutter.attach_tree(self.tree)
cutter.choose_leaves(leaves)
result = cutter.cut()
compare_nodes(self, result.get_newick()[0], expected_tree[0])
def test_cut_sub_branch(self):
# cut {E F}
leaves = {'A', 'B', 'C', 'D'}
expected_tree = loads("({A},({B},{C}){B C},{D}){A B C D};")
cutter = TreeCutter()
cutter.attach_tree(self.tree)
cutter.choose_leaves(leaves)
result = cutter.cut()
compare_nodes(self, result.get_newick()[0], expected_tree[0])
def test_cut_whole_branch(self):
# cut {B C}
leaves = {'A', 'D', 'E', 'F'}
expected_tree = loads("({A},({D},({E},{F}){E F}){D E F}){A D E F};")
cutter = TreeCutter()
cutter.attach_tree(self.tree)
cutter.choose_leaves(leaves)
result = cutter.cut()
compare_nodes(self, result.get_newick()[0], expected_tree[0])
def test_cut_leaf_from_second_layer(self):
# cut D
leaves = {'A', 'B', 'C', 'E', 'F'}
expected_tree = loads("({A},({B},{C}){B C},({E},{F}){E F}){A B C E F};")
cutter = TreeCutter()
cutter.attach_tree(self.tree)
cutter.choose_leaves(leaves)
result = cutter.cut()
compare_nodes(self, result.get_newick()[0], expected_tree[0])
def test_cut_leaf_from_two_parts_group(self):
# cut E
leaves = {'A', 'B', 'C', 'D', 'F'}
expected_tree = loads("({A},({B},{C}){B C},({D},{F}){D F}){A B C D F};")
cutter = TreeCutter()
cutter.attach_tree(self.tree)
cutter.choose_leaves(leaves)
result = cutter.cut()
compare_nodes(self, result.get_newick()[0], expected_tree[0])
def cut(file_name: str):
if not os.path.isfile(file_name):
print("Given file has to exist.")
return
tree = PhyTree()
tree.load_file(file_name)
wanted_leaves = set()
choose_leaves = True
while choose_leaves:
print(f"Chosen leaves: {wanted_leaves}")
leaf = input("Add leaf: ")
if leaf not in tree.leaves:
print(f"{leaf} was not among the tree's leaves")
continue
if leaf in wanted_leaves:
print(f"{leaf} is already among wanted leaves")
continue
wanted_leaves.add(leaf)
choose_leaves = not ask_if_finished()
cutter = TreeCutter()
cutter.attach_tree(tree)
cutter.choose_leaves(wanted_leaves)
new_tree = cutter.cut()
new_tree.save(file_name)