Merge pull request #248 from kayjan/fix-type-hints

Fix type hints

Author: kayjan

CHANGELOG.md

@@ -6,6 +6,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 ### Changed:
+- Binary Tree Constructor: Type hints to return more generic TypeVar for use with subclasses.
+- DAG Constructor: Type hints to return more generic TypeVar for use with subclasses.
+- [#247] Tree Construct: Type hints to return more generic TypeVar for use with subclasses.
+- Tree Modifier: Type hints to return more generic TypeVar for use with subclasses.
 - Misc: Documentation to include more contribution information and guidelines.
 
 ## [0.18.2] - 2024-06-01

bigtree/binarytree/construct.py

@@ -1,15 +1,17 @@
-from typing import List, Type
+from typing import List, Type, TypeVar
 
 from bigtree.node.binarynode import BinaryNode
+from bigtree.utils.assertions import assert_length_not_empty
 
 __all__ = ["list_to_binarytree"]
 
-from bigtree.utils.assertions import assert_length_not_empty
+T = TypeVar("T", bound=BinaryNode)
 
 
 def list_to_binarytree(
-    heapq_list: List[int], node_type: Type[BinaryNode] = BinaryNode
-) -> BinaryNode:
+    heapq_list: List[int],
+    node_type: Type[T] = BinaryNode,  # type: ignore[assignment]
+) -> T:
     """Construct tree from a list of numbers (int or float) in heapq format.
 
     Examples:
@@ -46,6 +48,6 @@ def list_to_binarytree(
     for idx, num in enumerate(heapq_list):
         if idx:
             parent_idx = int((idx + 1) / 2) - 1
-            node = node_type(num, parent=node_list[parent_idx])  # type: ignore
+            node = node_type(num, parent=node_list[parent_idx])
             node_list.append(node)
     return root_node

bigtree/dag/construct.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, Dict, List, Tuple, Type
+from typing import Any, Dict, List, Tuple, Type, TypeVar
 
 from bigtree.node.dagnode import DAGNode
 from bigtree.utils.assertions import (
@@ -20,11 +20,13 @@
 
 __all__ = ["list_to_dag", "dict_to_dag", "dataframe_to_dag"]
 
+T = TypeVar("T", bound=DAGNode)
+
 
 def list_to_dag(
     relations: List[Tuple[str, str]],
-    node_type: Type[DAGNode] = DAGNode,
-) -> DAGNode:
+    node_type: Type[T] = DAGNode,  # type: ignore[assignment]
+) -> T:
     """Construct DAG from list of tuples containing parent-child names.
     Note that node names must be unique.
 
@@ -44,8 +46,8 @@ def list_to_dag(
     """
     assert_length_not_empty(relations, "Input list", "relations")
 
-    node_dict: Dict[str, DAGNode] = dict()
-    parent_node = DAGNode()
+    node_dict: Dict[str, T] = dict()
+    parent_node: T = DAGNode()  # type: ignore[assignment]
 
     for parent_name, child_name in relations:
         if parent_name not in node_dict:
@@ -67,8 +69,8 @@ def list_to_dag(
 def dict_to_dag(
     relation_attrs: Dict[str, Any],
     parent_key: str = "parents",
-    node_type: Type[DAGNode] = DAGNode,
-) -> DAGNode:
+    node_type: Type[T] = DAGNode,  # type: ignore[assignment]
+) -> T:
     """Construct DAG from nested dictionary, ``key``: child name, ``value``: dictionary of parent names, attribute
     name, and attribute value.
     Note that node names must be unique.
@@ -97,8 +99,8 @@ def dict_to_dag(
     """
     assert_length_not_empty(relation_attrs, "Dictionary", "relation_attrs")
 
-    node_dict: Dict[str, DAGNode] = dict()
-    parent_node: DAGNode | None = None
+    node_dict: Dict[str, T] = dict()
+    parent_node: T | None = None
 
     for child_name, node_attrs in relation_attrs.items():
         node_attrs = node_attrs.copy()
@@ -133,8 +135,8 @@ def dataframe_to_dag(
     child_col: str = "",
     parent_col: str = "",
     attribute_cols: List[str] = [],
-    node_type: Type[DAGNode] = DAGNode,
-) -> DAGNode:
+    node_type: Type[T] = DAGNode,  # type: ignore[assignment]
+) -> T:
     """Construct DAG from pandas DataFrame.
     Note that node names must be unique.
 
@@ -200,8 +202,8 @@ def dataframe_to_dag(
     if sum(data[child_col].isnull()):
         raise ValueError(f"Child name cannot be empty, check column: {child_col}")
 
-    node_dict: Dict[str, DAGNode] = dict()
-    parent_node = DAGNode()
+    node_dict: Dict[str, T] = dict()
+    parent_node: T = DAGNode()  # type: ignore[assignment]
 
     for row in data.reset_index(drop=True).to_dict(orient="index").values():
         child_name = row[child_col]

bigtree/tree/construct.py

@@ -2,7 +2,7 @@
 
 import re
 from collections import OrderedDict, defaultdict
-from typing import Any, Dict, List, Optional, Tuple, Type
+from typing import Any, Dict, List, Optional, Tuple, Type, TypeVar
 
 from bigtree.node.node import Node
 from bigtree.tree.search import find_child_by_name, find_name
@@ -51,14 +51,16 @@
     "newick_to_tree",
 ]
 
+T = TypeVar("T", bound=Node)
+
 
 def add_path_to_tree(
-    tree: Node,
+    tree: T,
     path: str,
     sep: str = "/",
     duplicate_name_allowed: bool = True,
     node_attrs: Dict[str, Any] = {},
-) -> Node:
+) -> T:
     """Add nodes and attributes to existing tree *in-place*, return node of path added.
     Adds to existing tree from list of path strings.
 
@@ -136,11 +138,11 @@ def add_path_to_tree(
 
 
 def add_dict_to_tree_by_path(
-    tree: Node,
+    tree: T,
     path_attrs: Dict[str, Dict[str, Any]],
     sep: str = "/",
     duplicate_name_allowed: bool = True,
-) -> Node:
+) -> T:
     """Add nodes and attributes to tree *in-place*, return root of tree.
     Adds to existing tree from nested dictionary, ``key``: path, ``value``: dict of attribute name and attribute value.
 
@@ -208,7 +210,7 @@ def add_dict_to_tree_by_path(
     return root_node
 
 
-def add_dict_to_tree_by_name(tree: Node, name_attrs: Dict[str, Dict[str, Any]]) -> Node:
+def add_dict_to_tree_by_name(tree: T, name_attrs: Dict[str, Dict[str, Any]]) -> T:
     """Add attributes to existing tree *in-place*.
     Adds to existing tree from nested dictionary, ``key``: name, ``value``: dict of attribute name and attribute value.
 
@@ -254,13 +256,13 @@ def add_dict_to_tree_by_name(tree: Node, name_attrs: Dict[str, Dict[str, Any]])
 
 
 def add_dataframe_to_tree_by_path(
-    tree: Node,
+    tree: T,
     data: pd.DataFrame,
     path_col: str = "",
     attribute_cols: List[str] = [],
     sep: str = "/",
     duplicate_name_allowed: bool = True,
-) -> Node:
+) -> T:
     """Add nodes and attributes to tree *in-place*, return root of tree.
     Adds to existing tree from pandas DataFrame.
 
@@ -350,11 +352,11 @@ def add_dataframe_to_tree_by_path(
 
 
 def add_dataframe_to_tree_by_name(
-    tree: Node,
+    tree: T,
     data: pd.DataFrame,
     name_col: str = "",
     attribute_cols: List[str] = [],
-) -> Node:
+) -> T:
     """Add attributes to existing tree *in-place*.
     Adds to existing tree from pandas DataFrame.
 
@@ -418,13 +420,13 @@ def add_dataframe_to_tree_by_name(
 
 
 def add_polars_to_tree_by_path(
-    tree: Node,
+    tree: T,
     data: pl.DataFrame,
     path_col: str = "",
     attribute_cols: List[str] = [],
     sep: str = "/",
     duplicate_name_allowed: bool = True,
-) -> Node:
+) -> T:
     """Add nodes and attributes to tree *in-place*, return root of tree.
     Adds to existing tree from polars DataFrame.
 
@@ -516,11 +518,11 @@ def add_polars_to_tree_by_path(
 
 
 def add_polars_to_tree_by_name(
-    tree: Node,
+    tree: T,
     data: pl.DataFrame,
     name_col: str = "",
     attribute_cols: List[str] = [],
-) -> Node:
+) -> T:
     """Add attributes to existing tree *in-place*.
     Adds to existing tree from polars DataFrame.
 
@@ -586,8 +588,8 @@ def add_polars_to_tree_by_name(
 def str_to_tree(
     tree_string: str,
     tree_prefix_list: List[str] = [],
-    node_type: Type[Node] = Node,
-) -> Node:
+    node_type: Type[T] = Node,  # type: ignore[assignment]
+) -> T:
     r"""Construct tree from tree string
 
     Examples:
@@ -655,8 +657,8 @@ def list_to_tree(
     paths: List[str],
     sep: str = "/",
     duplicate_name_allowed: bool = True,
-    node_type: Type[Node] = Node,
-) -> Node:
+    node_type: Type[T] = Node,  # type: ignore[assignment]
+) -> T:
     """Construct tree from list of path strings.
 
     Path should contain ``Node`` name, separated by `sep`.
@@ -715,8 +717,8 @@ def list_to_tree(
 def list_to_tree_by_relation(
     relations: List[Tuple[str, str]],
     allow_duplicates: bool = False,
-    node_type: Type[Node] = Node,
-) -> Node:
+    node_type: Type[T] = Node,  # type: ignore[assignment]
+) -> T:
     """Construct tree from list of tuple containing parent-child names.
 
     Root node is inferred when parent is empty, or when name appears as parent but not as child.
@@ -764,8 +766,8 @@ def dict_to_tree(
     path_attrs: Dict[str, Any],
     sep: str = "/",
     duplicate_name_allowed: bool = True,
-    node_type: Type[Node] = Node,
-) -> Node:
+    node_type: Type[T] = Node,  # type: ignore[assignment]
+) -> T:
     """Construct tree from nested dictionary using path,
     ``key``: path, ``value``: dict of attribute name and attribute value.
 
@@ -854,8 +856,8 @@ def nested_dict_to_tree(
     node_attrs: Dict[str, Any],
     name_key: str = "name",
     child_key: str = "children",
-    node_type: Type[Node] = Node,
-) -> Node:
+    node_type: Type[T] = Node,  # type: ignore[assignment]
+) -> T:
     """Construct tree from nested recursive dictionary.
 
     - ``key``: `name_key`, `child_key`, or any attributes key.
@@ -901,8 +903,8 @@ def nested_dict_to_tree(
     assert_length_not_empty(node_attrs, "Dictionary", "node_attrs")
 
     def _recursive_add_child(
-        child_dict: Dict[str, Any], parent_node: Optional[Node] = None
-    ) -> Node:
+        child_dict: Dict[str, Any], parent_node: Optional[T] = None
+    ) -> T:
         """Recursively add child to tree, given child attributes and parent node.
 
         Args:
@@ -934,8 +936,8 @@ def dataframe_to_tree(
     attribute_cols: List[str] = [],
     sep: str = "/",
     duplicate_name_allowed: bool = True,
-    node_type: Type[Node] = Node,
-) -> Node:
+    node_type: Type[T] = Node,  # type: ignore[assignment]
+) -> T:
     """Construct tree from pandas DataFrame using path, return root of tree.
 
     `path_col` and `attribute_cols` specify columns for node path and attributes to construct tree.
@@ -1040,8 +1042,8 @@ def dataframe_to_tree_by_relation(
     parent_col: str = "",
     attribute_cols: List[str] = [],
     allow_duplicates: bool = False,
-    node_type: Type[Node] = Node,
-) -> Node:
+    node_type: Type[T] = Node,  # type: ignore[assignment]
+) -> T:
     """Construct tree from pandas DataFrame using parent and child names, return root of tree.
 
     Root node is inferred when parent name is empty, or when name appears in parent column but not in child column.
@@ -1138,7 +1140,7 @@ def _retrieve_attr(_row: Dict[str, Any]) -> Dict[str, Any]:
         node_attrs["name"] = _row[child_col]
         return node_attrs
 
-    def _recursive_add_child(parent_node: Node) -> None:
+    def _recursive_add_child(parent_node: T) -> None:
         """Recursive add child to tree, given current node.
 
         Args:
@@ -1168,8 +1170,8 @@ def polars_to_tree(
     attribute_cols: List[str] = [],
     sep: str = "/",
     duplicate_name_allowed: bool = True,
-    node_type: Type[Node] = Node,
-) -> Node:
+    node_type: Type[T] = Node,  # type: ignore[assignment]
+) -> T:
     """Construct tree from polars DataFrame using path, return root of tree.
 
     `path_col` and `attribute_cols` specify columns for node path and attributes to construct tree.
@@ -1275,8 +1277,8 @@ def polars_to_tree_by_relation(
     parent_col: str = "",
     attribute_cols: List[str] = [],
     allow_duplicates: bool = False,
-    node_type: Type[Node] = Node,
-) -> Node:
+    node_type: Type[T] = Node,  # type: ignore[assignment]
+) -> T:
     """Construct tree from polars DataFrame using parent and child names, return root of tree.
 
     Root node is inferred when parent name is empty, or when name appears in parent column but not in child column.
@@ -1373,7 +1375,7 @@ def _retrieve_attr(_row: Dict[str, Any]) -> Dict[str, Any]:
         node_attrs["name"] = _row[child_col]
         return node_attrs
 
-    def _recursive_add_child(parent_node: Node) -> None:
+    def _recursive_add_child(parent_node: T) -> None:
         """Recursive add child to tree, given current node.
 
         Args:
@@ -1402,8 +1404,8 @@ def newick_to_tree(
     tree_string: str,
     length_attr: str = "length",
     attr_prefix: str = "&&NHX:",
-    node_type: Type[Node] = Node,
-) -> Node:
+    node_type: Type[T] = Node,  # type: ignore[assignment]
+) -> T:
     """Construct tree from Newick notation, return root of tree.
 
     In the Newick Notation (or New Hampshire Notation)
@@ -1460,24 +1462,24 @@ def newick_to_tree(
     assert_length_not_empty(tree_string, "Tree string", "tree_string")
 
     # Store results (for tracking)
-    depth_nodes: Dict[int, List[Node]] = defaultdict(list)
+    depth_nodes: Dict[int, List[T]] = defaultdict(list)
     unlabelled_node_counter: int = 0
     current_depth: int = 1
     tree_string_idx: int = 0
 
     # Store states (for assertions and checks)
     current_state: NewickState = NewickState.PARSE_STRING
-    current_node: Optional[Node] = None
+    current_node: Optional[T] = None
     cumulative_string: str = ""
     cumulative_string_value: str = ""
 
     def _create_node(
-        _new_node: Optional[Node],
+        _new_node: Optional[T],
         _cumulative_string: str,
         _unlabelled_node_counter: int,
-        _depth_nodes: Dict[int, List[Node]],
+        _depth_nodes: Dict[int, List[T]],
         _current_depth: int,
-    ) -> Tuple[Node, int]:
+    ) -> Tuple[T, int]:
         """Create node at checkpoint.
 
         Args: