From c8985c48eebbd8a6a1b561add5d58cabe44e838f Mon Sep 17 00:00:00 2001
From: JLSteenwyk <jlsteenwyk@gmail.com>
Date: Sat, 14 Sep 2024 13:13:01 -0700
Subject: [PATCH] prune tree function now takes advantage of list comprehension

---
 phykit/services/tree/prune_tree.py | 23 +++++++++--------------
 1 file changed, 9 insertions(+), 14 deletions(-)

diff --git a/phykit/services/tree/prune_tree.py b/phykit/services/tree/prune_tree.py
index 947d68f..fd5bf3a 100644
--- a/phykit/services/tree/prune_tree.py
+++ b/phykit/services/tree/prune_tree.py
@@ -1,3 +1,5 @@
+from typing import Dict
+
 from .base import Tree
 
 from ...helpers.files import read_single_column_file_to_list
@@ -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,