Merge pull request #159 from kayjan/hyield-add-arg

Hyield add arg

Author: kayjan

bigtree/tree/construct.py

@@ -1182,10 +1182,14 @@ def _raise_value_error(tree_idx: int) -> None:
                 NewickState.PARSE_STRING,
                 NewickState.PARSE_ATTRIBUTE_NAME,
             ]:
+                if cumulative_string:
+                    _raise_value_error(tree_string_idx)
                 cumulative_string = tree_string[
                     tree_string_idx + 1 : quote_end_idx  # noqa: E203
                 ]
             else:
+                if cumulative_string_value:
+                    _raise_value_error(tree_string_idx)
                 cumulative_string_value = tree_string[
                     tree_string_idx + 1 : quote_end_idx  # noqa: E203
                 ]

bigtree/tree/export.py

@@ -377,6 +377,7 @@ def hprint_tree(
     tree: T,
     node_name_or_path: str = "",
     max_depth: int = 0,
+    intermediate_node_name: bool = True,
     style: str = "const",
     custom_style: Iterable[str] = [],
 ) -> None:
@@ -457,12 +458,14 @@ def hprint_tree(
         tree (Node): tree to print
         node_name_or_path (str): node to print from, becomes the root node of printing
         max_depth (int): maximum depth of tree to print, based on `depth` attribute, optional
+        intermediate_node_name (bool): indicator if intermediate nodes have node names, defaults to True
         style (str): style of print, defaults to const style
         custom_style (Iterable[str]): style of icons, used when `style` is set to 'custom'
     """
     result = hyield_tree(
         tree,
         node_name_or_path=node_name_or_path,
+        intermediate_node_name=intermediate_node_name,
         max_depth=max_depth,
         style=style,
         custom_style=custom_style,
@@ -474,6 +477,7 @@ def hyield_tree(
     tree: T,
     node_name_or_path: str = "",
     max_depth: int = 0,
+    intermediate_node_name: bool = True,
     style: str = "const",
     custom_style: Iterable[str] = [],
 ) -> List[str]:
@@ -555,6 +559,7 @@ def hyield_tree(
         tree (Node): tree to print
         node_name_or_path (str): node to print from, becomes the root node of printing
         max_depth (int): maximum depth of tree to print, based on `depth` attribute, optional
+        intermediate_node_name (bool): indicator if intermediate nodes have node names, defaults to True
         style (str): style of print, defaults to const style
         custom_style (Iterable[str]): style of icons, used when `style` is set to 'custom'
 
@@ -609,9 +614,10 @@ def hyield_tree(
 
     # Calculate padding
     space = " "
-    padding_depths = {}
-    for _idx, _children in enumerate(levelordergroup_iter(tree)):
-        padding_depths[_idx + 1] = max([len(node.node_name) for node in _children])
+    padding_depths = collections.defaultdict(int)
+    if intermediate_node_name:
+        for _idx, _children in enumerate(levelordergroup_iter(tree)):
+            padding_depths[_idx + 1] = max([len(node.node_name) for node in _children])
 
     def _hprint_branch(_node: T) -> Tuple[List[str], int]:
         """Get string for tree horizontally.
@@ -623,17 +629,19 @@ def _hprint_branch(_node: T) -> Tuple[List[str], int]:
         Returns:
             (Tuple[List[str], int]): Intermediate/final result for node, index of branch
         """
-        padding_depth = padding_depths[_node.depth]
-        padding = space * (padding_depth + 4)
-        node_name_centered = _node.node_name.center(padding_depth)
+        node_name_centered = _node.node_name.center(padding_depths[_node.depth])
 
         children = list(_node.children)
         if not len(children):
             node_str = f"{style_branch} {node_name_centered.rstrip()}"
             return [node_str], 0
 
         result, result_nrow, result_idx = [], [], []
-        node_str = f"""{style_branch} {node_name_centered} {style_branch}"""
+        if intermediate_node_name:
+            node_str = f"""{style_branch} {node_name_centered} {style_branch}"""
+        else:
+            node_str = f"""{style_branch}{style_branch}{style_branch}"""
+        padding = space * len(node_str)
         for idx, child in enumerate(children):
             result_child, result_branch_idx = _hprint_branch(child)
             result.extend(result_child)

tests/tree/test_construct.py

@@ -2554,6 +2554,24 @@ def test_newick_to_tree_node_type(self):
         )
         assert_tree_structure_node_root(root)
 
+    def test_newick_to_tree_invalid_character_error(self):
+        newick_strs_error = [
+            (
+                """((d,(g,h)e)b,(f)c)"'a'\"""",
+                19,
+            ),  # NewickCharacter.ATTR_QUOTE, wrong order of bracket (name)
+            (
+                """((d,(g,h)e)b,(f[age="'38'"])c)a""",
+                21,
+            ),  # NewickCharacter.ATTR_QUOTE, wrong order of bracket (attr value)
+        ]
+        for newick_str, error_idx in newick_strs_error:
+            with pytest.raises(ValueError) as exc_info:
+                newick_to_tree(newick_str)
+            assert str(exc_info.value) == Constants.ERROR_NODE_NEWICK_NOT_CLOSED.format(
+                index=error_idx
+            )
+
 
 def assert_tree_structure_phylogenetic(root):
     assert root.max_depth == 4, f"Expected max_depth 4, received {root.max_depth}"

tests/tree/test_export.py

@@ -316,6 +316,30 @@ def test_hprint_tree_child_node_path(tree_node):
             node_name_or_path="a/b",
         )
 
+    @staticmethod
+    def test_hprint_tree_intermediate_node_name(tree_node):
+        expected_str = (
+            "       ┌─ d\n   ┌───┤   ┌─ g\n───┤   └───┤\n   │       └─ h\n   └───── f\n"
+        )
+        assert_print_statement(
+            hprint_tree,
+            expected_str,
+            tree=tree_node,
+            intermediate_node_name=False,
+        )
+
+    @staticmethod
+    def test_hprint_tree_intermediate_node_name_diff_node_name_length(tree_node):
+        tree_node["b"].name = "bcde"
+        tree_node["c"]["f"].name = "fghijk"
+        expected_str = "       ┌─ d\n   ┌───┤   ┌─ g\n───┤   └───┤\n   │       └─ h\n   └───── fghijk\n"
+        assert_print_statement(
+            hprint_tree,
+            expected_str,
+            tree=tree_node,
+            intermediate_node_name=False,
+        )
+
     @staticmethod
     def test_hprint_tree_style_ansi(tree_node):
         expected_str = "           /- d\n     /- b -+     /- g\n- a -+     \\- e -+\n     |           \\- h\n     \\- c --- f\n"