Repository: cogent3/cogent3
Pull request #2372 (2025-06-23 22:47:33)

Refactor work on tree class methods
## Summary by Sourcery

Refactor Tree and TreeBuilder methods to streamline signatures, add static type hints, and improve subtree extraction and parameter merging logic; update Newick parser to accept missing semicolons and adjust tests to validate new behavior.

Bug Fixes:
- Handle Newick strings lacking trailing semicolons in DndParser

Enhancements:
- Introduce is_number helper and generalize prune to sum numeric parameters across nodes
- Replace legacy _get_sub_tree with a clearer get_sub_tree implementation that preserves unrooted status and supports tip-only extraction
- Simplify method signatures (copy, unrooted_deepcopy, _default_tree_constructor) and add type annotations throughout Tree and TreeBuilder classes

Tests:
- Revise and add tests for get_sub_tree behavior, param merging, unrooted cases, and parsing named root nodes without semicolons


Patch:
diff --git a/src/cogent3/core/tree.py b/src/cogent3/core/tree.py
index c9f9b15ea..5020b3ac0 100644
--- a/src/cogent3/core/tree.py
+++ b/src/cogent3/core/tree.py
@@ -29,7 +29,6 @@ from __future__ import annotations
 
 import contextlib
 import json
-import numbers
 import re
 import typing
 from copy import deepcopy
@@ -39,7 +38,7 @@ from operator import or_
 from random import choice, shuffle
 
 import typing_extensions
-from numpy import argsort, ceil, log, zeros
+from numpy import argsort, ceil, floating, integer, log, zeros
 
 from cogent3._version import __version__
 from cogent3.maths.stats.test import correlation
@@ -57,6 +56,10 @@ if typing_extensions.TYPE_CHECKING:
     from cogent3.evolve.fast_distance import DistanceMatrix
 
 
+def is_number(v) -> bool:
+    return isinstance(v, (int, float, integer, floating))
+
+
 def distance_from_r_squared(m1, m2):
     """Estimates distance as 1-r^2: no correl = max distance"""
     return 1 - (correlation(m1.flat, m2.flat)[0]) ** 2
@@ -292,11 +295,9 @@ class TreeNode:
         return len(self.children)
 
     # support for copy module
-    def copy(self, memo=None, _nil=None, constructor="ignored"):
+    def copy(self, memo=None, _nil=None, constructor: None = None):
         """Returns a copy of self using an iterative approach"""
 
-        _nil = _nil or []
-
         root = _copy_node(self)
         nodes_stack = [[root, self, len(self.children)]]
 
@@ -486,7 +487,7 @@ class TreeNode:
         # we put tips on the right
         right_name = edge_name if is_tip else f"{edge_name}-R"
         left_name = f"{edge_name}-root" if is_tip else f"{edge_name}-L"
-        length = getattr(node, "length", 0.0) / 2
+        length = (getattr(node, "length", 0.0) or 0) / 2
         parent = node.parent
         parent.children.remove(node)
         node.parent = None
@@ -712,7 +713,7 @@ class TreeNode:
                     result[i, j] = 1
         return result, node_list
 
-    def _default_tree_constructor(self):
+    def _default_tree_constructor(self) -> typing.Callable[[T, list[T]], T]:
         return TreeBuilder(constructor=self.__class__).edge_from_edge
 
     def name_unnamed_nodes(self) -> None:
@@ -809,7 +810,11 @@ class TreeNode:
                 child.parent = curr_parent
                 if not has_length:
                     continue
-                child.length = (child.length or 0.0) + (node.length or 0.0)
+                for key, value in child.params.items():
+                    node_val = node.params.get(key)
+                    if is_number(value) or is_number(node_val):
+                        value = (value or 0.0) + (node_val or 0.0)
+                    child.params[key] = value
 
     def same_shape(self, other):
         """Ignores lengths and order, so trees should be sorted first"""
@@ -1005,74 +1010,6 @@ class TreeNode:
             if c is not None and c is not parent
         ]
 
-    def _get_sub_tree(
-        self,
-        included_names: list[str],
-        constructor: TreeNode | PhyloNode | None = None,
-        keep_root: bool = False,
-        tips_only: bool = False,
-    ) -> typing_extensions.Self | None:
-        """An equivalent node with possibly fewer children, or None"""
-
-        # Renumber autonamed edges
-        if constructor is None:
-            constructor = self._default_tree_constructor()
-
-        if (self.name in included_names and not tips_only) or (
-            self.name in included_names and self.istip()
-        ):
-            return self.deepcopy(constructor=constructor)
-        # don't need to pass keep_root to children, though
-        # internal nodes will be elminated this way
-        children = []
-        for child in self.children:
-            st = child._get_sub_tree(included_names, constructor, tips_only=tips_only)
-            children.append(st)
-
-        children = [child for child in children if child is not None]
-        if len(children) == 0:
-            result = None
-        elif len(children) == 1 and not keep_root:
-            # Merge parameter dictionaries by adding lengths and making
-            # weighted averages of other parameters.  This should probably
-            # be moved out of here into a ParameterSet class (Model?) or
-            # tree subclass.
-            params = {}
-            child = children[0]
-            if self.length is not None and child.length is not None:
-                shared_params = [
-                    n
-                    for (n, v) in list(self.params.items())
-                    if v is not None
-                    and child.params.get(n) is not None
-                    and n != "length"
-                ]
-                length = self.length + child.length
-                if length:
-                    params = {}
-                    for n in shared_params:
-                        self_val = self.params[n]
-                        child_val = child.params[n]
-                        is_scalar = True
-                        for i in (self_val, child_val):
-                            if not isinstance(i, numbers.Number):
-                                is_scalar = False
-                                break
-                        if is_scalar:
-                            val = (
-                                self_val * self.length + child_val * child.length
-                            ) / length
-                        else:
-                            val = self_val
-                        params[n] = val
-
-                    params["length"] = length
-            result = child
-            result.params = params
-        else:
-            result = constructor(self, tuple(children))
-        return result
-
     @c3warn.deprecated_args(
         "2025.6",
         "consistency between methods",
@@ -1101,29 +1038,69 @@ class TreeNode:
         tips_only
             only tip names matching name_list are allowed
 
+        Notes
+        -----
+        An unrooted tree returns an unrooted subtree.
         """
-        edge_names = set(self.get_node_names(include_self=True, tips_only=tips_only))
-        if not ignore_missing:
-            # this may take a long time
-            for name in name_list:
-                if name not in edge_names:
-                    msg = f"edge {name!r} not found in tree"
-                    raise ValueError(msg)
-
-        new_tree = self._get_sub_tree(
-            name_list, keep_root=keep_root, tips_only=tips_only
-        )
-        if new_tree is None:
-            msg = "no tree created in make sub tree"
-            raise TreeError(msg)
-        if new_tree.istip():
-            msg = "only a tip was returned from selecting sub tree"
-            raise TreeError(msg)
-        new_tree.name = "root"
-        # keep unrooted
-        if len(self.children) > 2:
-            new_tree = new_tree.unrooted()
-        return new_tree
+        # find all the selected nodes
+        allowed = set(name_list)
+        old_nodes = {}
+        found = set()
+        for old_node in self.preorder(include_self=True):
+            if old_node.name not in allowed:
+                continue
+
+            found.add(old_node.name)
+            old_nodes[id(old_node)] = old_node
+            # find all nodes connecting required nodes to root,
+            # skipping if already present
+            parent = old_node.parent
+            while parent is not None and (parent_id := id(parent)) not in old_nodes:
+                old_nodes[parent_id] = parent
+                parent = parent.parent
+
+            if not tips_only and not old_node.is_tip():
+                # add all descendant nodes too
+                for n in old_node.preorder():
+                    old_nodes[id(n)] = n
+
+        if found != allowed and not ignore_missing:
+            msg = f"edges {allowed - found} not found in tree"
+            raise ValueError(msg)
+
+        # make new nodes and also map old id's to new id's
+        make_node = self.__class__
+        self_2_new = {}
+        new_nodes = {}
+        for self_id, old_node in old_nodes.items():
+            new_node = make_node(old_node.name, params={**old_node.params})
+            new_nodes[id(new_node)] = new_node
+            self_2_new[self_id] = id(new_node)
+
+        # connect the nodes
+        for self_id, old_node in old_nodes.items():
+            if old_node.parent is None:
+                continue
+
+            new_node = new_nodes[self_2_new[self_id]]
+            new_parent_id = self_2_new[id(old_node.parent)]
+            # the following assignment also adds the new_node as
+            # a child to parent
+            new_node.parent = new_nodes[new_parent_id]
+
+        result_root = new_nodes[self_2_new[id(self)]]
+        if not keep_root and len(result_root.children) == 1:
+            while len(result_root.children) == 1:
+                result_root = result_root.children[0]
+
+            result_root.parent = None
+
+        result_root.prune()
+        if len(self.children) > 2 and len(result_root.children) == 2:
+            # we preserve the "unrooted" status in the result tree
+            result_root = result_root.unrooted()
+        result_root.name = "root"
+        return result_root
 
     def _edgecount(self, parent, cache):
         """ "The number of edges beyond 'parent' in the direction of 'self',
@@ -1670,15 +1647,19 @@ class TreeNode:
             dists[tip.name] = cum_sum
         return dists
 
-    def same_topology(self, other):
+    def same_topology(self, other) -> bool:
         """Tests whether two trees have the same topology."""
         tip_names = self.get_tip_names()
         root_at = tip_names[0]
-        me = self.rooted_with_tip(root_at).sorted(tip_names)
-        them = other.rooted_with_tip(root_at).sorted(tip_names)
+        me = self.rooted(root_at).sorted(tip_names)
+        them = other.rooted(root_at).sorted(tip_names)
         return self is other or me.same_shape(them)
 
-    def unrooted_deepcopy(self, constructor=None, parent=None):
+    def unrooted_deepcopy(
+        self,
+        constructor: type | None = None,
+        parent: typing_extensions.Self | None = None,
+    ) -> typing_extensions.Self:
         """
         Returns a deepcopy of the tree using unrooted traversal.
 
@@ -2219,39 +2200,44 @@ def split_name_and_support(name_field: str | None) -> tuple[str | None, float |
     return name, support_value
 
 
+T = typing.TypeVar("T", "TreeNode", "PhyloNode")
+
+
 class TreeBuilder:
     # Some tree code which isn't needed once the tree is finished.
     # Mostly exists to give edges unique names
     # children must be created before their parents.
 
-    def __init__(self, mutable=False, constructor=PhyloNode) -> None:
+    def __init__(self, mutable: bool = False, constructor: type = PhyloNode) -> None:
         self._used_names = {"edge": -1}
         self._known_edges = {}
         self.TreeNodeClass = constructor
 
-    def _unique_name(self, name):
+    def _unique_name(self, name: str) -> str:
         # Unnamed edges become edge.0, edge.1 edge.2 ...
         # Other duplicates go mouse mouse.2 mouse.3 ...
         if not name:
             name = "edge"
         if name in self._used_names:
             self._used_names[name] += 1
-            name += "." + str(self._used_names[name])
+            name += f".{self._used_names[name]!s}"
             # in case of names like 'edge.1.1'
             name = self._unique_name(name)
         else:
             self._used_names[name] = 1
         return name
 
-    def _params_for_edge(self, edge):
+    def _params_for_edge(self, edge: T) -> dict:
         # default is just to keep it
         return edge.params
 
-    def edge_from_edge(self, edge, children, params=None):
+    def edge_from_edge(
+        self, edge: T | None, children: list[T], params: dict | None = None
+    ) -> T:
         """Callback for tree-to-tree transforms like get_sub_tree"""
         if edge is None:
             assert not params
-            return self.create_edge(children, "root", {}, False)
+            return self.create_edge(children, "root", {}, name_loaded=False)
         if params is None:
             params = self._params_for_edge(edge)
         return self.create_edge(
@@ -2261,7 +2247,9 @@ class TreeBuilder:
             name_loaded=edge.name_loaded,
         )
 
-    def create_edge(self, children, name, params, name_loaded=True):
+    def create_edge(
+        self, children: list[T], name: str, params: dict, name_loaded: bool = True
+    ) -> T:
         """Callback for newick parser"""
         if children is None:
             children = []
@@ -2285,7 +2273,7 @@ def make_tree(
     tip_names: list[str] | None = None,
     format: str | None = None,
     underscore_unmunge: bool = False,
-    source: str | None = None,
+    source: str | pathlib.Path | None = None,
 ) -> PhyloNode | TreeNode:
     """Initialises a tree.
 
@@ -2368,8 +2356,8 @@ def load_tree(
     -------
     PhyloNode
     """
-    file_format, _ = get_format_suffixes(filename)
-    format = format or file_format
+    fmt, _ = get_format_suffixes(filename)
+    format = format or fmt
     if format == "json":
         tree = load_from_json(filename, (TreeNode, PhyloNode))
         tree.source = str(filename)
diff --git a/src/cogent3/parse/tree.py b/src/cogent3/parse/tree.py
index 8f05cfd5e..aefcea2ea 100644
--- a/src/cogent3/parse/tree.py
+++ b/src/cogent3/parse/tree.py
@@ -86,6 +86,10 @@ def DndParser(lines, constructor=PhyloNode, unescape_name=False):
     explicitly).
     """
     data = lines if isinstance(lines, str) else "".join(lines)
+    data = data.strip()
+    if not data.endswith(";"):
+        data = f"{data};"
+
     # skip arb comment stuff if present: start at first paren
     paren_index = data.find("(")
     data = data[paren_index:]


Test patch:
diff --git a/tests/test_core/test_tree.py b/tests/test_core/test_tree.py
index dc235a3d8..68fdf0011 100644
--- a/tests/test_core/test_tree.py
+++ b/tests/test_core/test_tree.py
@@ -23,7 +23,7 @@ def test_make_tree():
     # NOTE: This method now sits in cogent3.__init__
 
     t_str = "(a_a:10,(b_b:2,c_c:4):5);"
-    # NOTE: Tree quotes these labels because they have underscores in them.
+    # Tree quotes these labels because they have underscores in them.
     result_str = "('a_a':10.0,('b_b':2.0,'c_c':4.0):5.0);"
     t = make_tree(treestring=t_str)
     names = [i.name for i in t.tips()]
@@ -45,15 +45,6 @@ def test_make_tree():
     assert t.get_tip_names() == ["1", "2", "3"]
 
 
-def _new_child(old_node, constructor):
-    """Returns new_node which has old_node as its parent."""
-    new_node = constructor()
-    new_node.parent = old_node
-    if old_node is not None and new_node not in old_node.children:
-        old_node.children.append(new_node)
-    return new_node
-
-
 tree_std = """\
         ((a:1, b:2, c:3)abc:0.1, (d:4, (e:5, f:6)ef:0.2)def:0.3);
 """
@@ -1929,15 +1920,20 @@ def test_balanced():
 
 
 def test_params_merge():
-    t = make_tree(treestring="((((a,b)ab,c)abc),d)")
-    for label, length, beta in [("a", 1, 20), ("b", 3, 2.0), ("ab", 4, 5.0)]:
-        t.get_node_matching_name(label).params = {"length": length, "beta": beta}
+    t = make_tree(treestring="((((a:1,b:3)ab:4,c)abc),d)")
+    for label, beta in [("a", 20), ("b", 2.0), ("ab", 5.0)]:
+        t.get_node_matching_name(label).params |= {"beta": beta}
     t = t.get_sub_tree(["b", "c", "d"])
+    # previous implementation on merge was
+    # dividing the sum of parameters across nodes by the lengths
+    # we no longer try and support this
     assert t.get_node_matching_name("b").params == {
         "length": 7,
-        "beta": float(2 * 3 + 4 * 5) / (3 + 4),
+        "beta": 2 + 5,
     }
-    assert str(t.get_sub_tree(["b", "c", "xxx"], ignore_missing=True)) == "(b:7,c);"
+    ex = make_tree("(b:7,c);")
+    st = t.get_sub_tree(["b", "c", "xxx"], ignore_missing=True)
+    assert ex.same_topology(st)
     with pytest.raises(ValueError):
         # should raise ValueError if a tip is not found
         t.get_sub_tree(["b", "c", "xxx"])
@@ -2069,12 +2065,18 @@ def test_getsubtree(request, tree_fxt, otu_fxt, nwk_fxt):
     tree = request.getfixturevalue(tree_fxt)
     otu_names = request.getfixturevalue(otu_fxt)
     newick_reduced = request.getfixturevalue(nwk_fxt)
-    subtree = tree.unrooted().get_sub_tree(otu_names)
+    tree = tree.unrooted()
+    subtree = tree.get_sub_tree(otu_names)
 
-    new_tree = make_tree(treestring=newick_reduced).unrooted()
+    exp_tree = make_tree(treestring=newick_reduced).unrooted()
 
-    # check we get the same names
-    assert str(subtree) == str(new_tree)
+    # names match
+    assert set(subtree.get_tip_names()) == set(exp_tree.get_tip_names())
+    # topology matches
+    assert exp_tree.same_topology(subtree)
+    # newick correctly formed means we reconstruct the same topology
+    from_nwk = make_tree(subtree.get_newick())
+    assert exp_tree.same_topology(from_nwk)
 
 
 def test_getsubtree_2():
@@ -2191,7 +2193,7 @@ def test_load_tree(tmp_path, tree_fxt, request):
     got = load_tree(tree_path)
     assert isinstance(got, PhyloNode)
     assert got.get_newick() == tree.get_newick()
-    assert got.get_node_names() == tree.get_node_names()
+    assert set(got.get_node_names()) == set(tree.get_node_names())
 
 
 @pytest.mark.parametrize("tree_fxt", ["tree_small", "tree_big"])
@@ -2203,7 +2205,7 @@ def test_load_tree_from_json(tmp_path, tree_fxt, request):
     got = load_tree(json_path)
     assert isinstance(got, PhyloNode)
     assert got.get_newick() == tree.get_newick()
-    assert got.get_node_names() == tree.get_node_names()
+    assert set(got.get_node_names()) == set(tree.get_node_names())
     # now try using non json suffix
     json_path = tmp_path / "tree.txt"
     tree.write(json_path, format="json")
@@ -2259,6 +2261,20 @@ def test_getsubtree_6():
     assert all(sub1.get_node_matching_name(name).params["non-scalar"] for name in names)
 
 
+def test_get_subtree_7():
+    t = make_tree(treestring="((((a:1,b:3)ab:4,c)abc),d)")
+    st = t.get_sub_tree(["b", "c"])
+    assert st.is_root()
+
+
+@pytest.mark.parametrize("treestring", ["(a,b,(c,(d,e)))", "((a,b),(c,(d,e)))"])
+def test_get_subtree_8(treestring):
+    # an unrooted tree returns an unrooted subtree
+    tree = make_tree(treestring=treestring)
+    st = tree.get_sub_tree(["c", "d", "e"])
+    assert len(st.children) == len(tree.children)
+
+
 def test_get_edge_names_big(tree_big, otu_names_big):
     """testing (well, exercising at least), getedgenames"""
     # Fell over on small tree because "stem descended from root
@@ -2499,6 +2515,13 @@ def test_dataset_tree_source():
     assert tr.source.endswith("murphy.tree")
 
 
+def test_rooted_is_rooted_diff_root_name():
+    treestring = "(c,b)abc"
+    tree = make_tree(treestring=treestring)
+    rooted = tree.rooted("b")
+    assert rooted.same_topology(tree)
+
+
 def test_rooted_is_rooted():
     treestring = "((Human,Chimpanzee),Gorilla)"
     tree = make_tree(treestring=treestring)
diff --git a/tests/test_parse/test_tree.py b/tests/test_parse/test_tree.py
index 8196f88d0..d60b196ee 100644
--- a/tests/test_parse/test_tree.py
+++ b/tests/test_parse/test_tree.py
@@ -2,6 +2,8 @@
 
 from unittest import TestCase
 
+import pytest
+
 from cogent3.core.tree import PhyloNode
 from cogent3.parse.tree import DndParser, DndTokenizer, RecordError
 
@@ -347,3 +349,12 @@ def test_make_tree_with_nhx(DATA_DIR):
     assert set(vert.params["other"]) == {'&sCF="50.68"', 'sDF1="24.42"'}
     vert = got.get_connecting_node("t5", "t6")
     assert set(vert.params["other"]) == {'&sCF="77.1"', 'sDF1="11.53"'}
+
+
+@pytest.mark.parametrize(
+    "treestring", ["(a,b)ab;", "(a,b)ab", "(a,b,(c,(d,e)))ab", "(a:3,b)ab"]
+)
+def test_parse_named_root_nodes(treestring):
+    tree = DndParser(treestring, PhyloNode)
+    assert tree.name == "ab"
+    assert tree.parent is None