From 3753f2f01893d0af9b860a298062dc40460adf31 Mon Sep 17 00:00:00 2001
From: JLSteenwyk <jlsteenwyk@gmail.com>
Date: Sun, 15 Sep 2024 11:44:54 -0700
Subject: [PATCH] added type hints and cleaned up the tbs function

---
 phykit/services/tree/terminal_branch_stats.py | 42 +++++++++++++------
 .../test_terminal_branch_stats_integration.py |  8 ++--
 2 files changed, 33 insertions(+), 17 deletions(-)

diff --git a/phykit/services/tree/terminal_branch_stats.py b/phykit/services/tree/terminal_branch_stats.py
index 253b577..2c69554 100644
--- a/phykit/services/tree/terminal_branch_stats.py
+++ b/phykit/services/tree/terminal_branch_stats.py
@@ -1,4 +1,7 @@
 import sys
+from typing import Dict, List, Tuple, Union
+
+from Bio.Phylo import Newick
 
 from .base import Tree
 
@@ -14,7 +17,8 @@ def __init__(self, args) -> None:
 
     def run(self):
         tree = self.read_tree_file()
-        _, stats, lengths_and_names = self.calculate_terminal_branch_stats(tree)
+        _, stats, lengths_and_names = \
+            self.calculate_terminal_branch_stats(tree)
 
         if self.verbose:
             try:
@@ -25,17 +29,25 @@ def run(self):
         else:
             print_summary_statistics(stats)
 
-    def process_args(self, args):
+    def process_args(self, args) -> Dict[str, str]:
         return dict(tree_file_path=args.tree, verbose=args.verbose)
 
-    def get_terminal_branch_lengths(self, tree) -> list:
+    def get_terminal_branch_lengths(
+        self,
+        tree: Newick.Tree,
+    ) -> List[
+        Union[
+            float,
+            List[List[Union[float, str]]],
+        ]
+    ]:
         """
         loop through tree and get all terminal branch lengths
         """
         terminal_branch_lengths = []
         lengths_and_names = []
         for terminal_branch in tree.get_terminals():
-            if terminal_branch.branch_length != None:
+            if terminal_branch.branch_length is not None:
                 temp = []
                 temp.append(terminal_branch.branch_length)
                 terminal_branch_lengths.append(terminal_branch.branch_length)
@@ -44,7 +56,10 @@ def get_terminal_branch_lengths(self, tree) -> list:
 
         return terminal_branch_lengths, lengths_and_names
 
-    def check_tree_has_branch_lengths(self, terminal_branch_lengths: list) -> None:
+    def check_tree_has_branch_lengths(
+        self,
+        terminal_branch_lengths: List[float],
+    ) -> None:
         """
         if tree has no branch lengths, exit
         """
@@ -54,16 +69,19 @@ def check_tree_has_branch_lengths(self, terminal_branch_lengths: list) -> None:
             )
             sys.exit()
 
-    def calculate_terminal_branch_stats(self, tree):
-        # save terminal branch lengths to terminal_branch_lengths
-        terminal_branch_lengths, lengths_and_names = self.get_terminal_branch_lengths(
-            tree
-        )
+    def calculate_terminal_branch_stats(
+        self,
+        tree: Newick.Tree,
+    ) -> Tuple[
+        List[float],
+        Dict[str, float],
+        List[List[Union[float, str]]],
+    ]:
+        terminal_branch_lengths, lengths_and_names = \
+            self.get_terminal_branch_lengths(tree)
 
-        # If the phylogeny had no branch lengths, inform user and quit
         self.check_tree_has_branch_lengths(terminal_branch_lengths)
 
-        # calculate summary stats
         stats = calculate_summary_statistics_from_arr(terminal_branch_lengths)
 
         return terminal_branch_lengths, stats, lengths_and_names
diff --git a/tests/integration/tree/test_terminal_branch_stats_integration.py b/tests/integration/tree/test_terminal_branch_stats_integration.py
index ff802b5..55bfb6a 100644
--- a/tests/integration/tree/test_terminal_branch_stats_integration.py
+++ b/tests/integration/tree/test_terminal_branch_stats_integration.py
@@ -1,9 +1,7 @@
-import pytest
-import sys
-from math import isclose
 from mock import patch, call
 from pathlib import Path
-from textwrap import dedent
+import pytest
+import sys
 
 from phykit.phykit import Phykit
 
@@ -132,4 +130,4 @@ def test_terminal_branch_stats_alias(self, mocked_print):
             Phykit()
 
         assert pytest_wrapped_e.type == SystemExit
-        assert pytest_wrapped_e.value.code == 2
\ No newline at end of file
+        assert pytest_wrapped_e.value.code == 2