Skip to content

Commit

Permalink
prune tree function now takes advantage of list comprehension
Browse files Browse the repository at this point in the history
  • Loading branch information
JLSteenwyk committed Sep 14, 2024
1 parent 4d6f8b1 commit c8985c4
Showing 1 changed file with 9 additions and 14 deletions.
23 changes: 9 additions & 14 deletions phykit/services/tree/prune_tree.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
from typing import Dict

from .base import Tree

from ...helpers.files import read_single_column_file_to_list
Expand All @@ -7,32 +9,25 @@ class PruneTree(Tree):
def __init__(self, args) -> None:
super().__init__(**self.process_args(args))

def run(self):
def run(self) -> None:
tree = self.read_tree_file()

taxa = read_single_column_file_to_list(self.list_of_taxa)

if self.keep:
tips_in_tree = []
for term in tree.get_terminals():
tips_in_tree.append(term.name)
tips_in_tree = [term.name for term in tree.get_terminals()]
taxa = [x for x in tips_in_tree if x not in taxa]

tree = self.prune_tree_using_taxa_list(tree, taxa)

self.write_tree_file(tree, self.output_file_path)

def process_args(self, args):
def process_args(self, args) -> Dict[str, str]:
tree_file_path = args.tree
if args.output is None:
output_file_path = f"{tree_file_path}.pruned"
else:
output_file_path = f"{args.output}"

if args.keep is None:
keep = True
else:
keep = args.keep
output_file_path = \
f"{args.output}" if args.output else f"{tree_file_path}.pruned"

keep = True if args.keep is None else args.keep

return dict(
tree_file_path=tree_file_path,
Expand Down

0 comments on commit c8985c4

Please sign in to comment.