Coverage for src/pyhrp/treelib.py: 100%

1"""A lightweight binary tree implementation to replace the binarytree dependency. 

2 

3This module provides a simple Node class that can be used to create binary trees. 

4It implements only the functionality needed by the pyhrp package. 

5""" 

6 

7from __future__ import annotations 

8 

9from collections import deque 

10from collections.abc import Iterator, Sequence 

11from typing import Generic, TypeVar 

12 

13# Type for node values 

14NodeValue = int | float | str 

15 

16T = TypeVar("T", bound=NodeValue) 

17 

18__all__ = ["Node"] 

19 

20 

21class Node(Generic[T]): 

22 """A binary tree node with left and right children. 

23 

24 This class implements the minimal functionality needed from the binarytree.Node class 

25 that is used in the pyhrp package. 

26 

27 Attributes: 

28 value: The value of the node 

29 left: The left child node 

30 right: The right child node 

31 """ 

32 

33 def __init__(self, value: T, left: Node[T] | None = None, right: Node[T] | None = None) -> None: 

34 """Initialize a new Node. 

35 

36 Args: 

37 value: The value of the node 

38 left: The left child node 

39 right: The right child node 

40 """ 

41 self.value = value 

42 self.left = left 

43 self.right = right 

44 

45 @property 

46 def is_leaf(self) -> bool: 

47 """Check if this node is a leaf node (has no children). 

48 

49 Returns: 

50 bool: True if this is a leaf node, False otherwise 

51 """ 

52 return self.left is None and self.right is None 

53 

54 @property 

55 def leaves(self) -> Sequence[Node[T]]: 

56 """Get all leaf nodes in the tree rooted at this node. 

57 

58 Returns: 

59 List[Node]: List of all leaf nodes 

60 """ 

61 # Iterative depth-first walk with an explicit stack: the leaf order is the 

62 # same left-to-right order the recursive form produced, but the traversal 

63 # depth is bounded by the heap rather than by sys.getrecursionlimit(). A 

64 # chain-degenerate cluster tree is as deep as it is wide, so recursing here 

65 # put a ceiling on the number of assets the package could handle. 

66 result: list[Node[T]] = [] 

67 stack: list[Node[T]] = [self] 

68 while stack: 

69 node = stack.pop() 

70 if node.is_leaf: 

71 result.append(node) 

72 continue 

73 # Right first, so the left subtree is popped and emitted first. 

74 if node.right is not None: 

75 stack.append(node.right) 

76 if node.left is not None: 

77 stack.append(node.left) 

78 

79 return result 

80 

81 @property 

82 def levels(self) -> list[list[Node[T]]]: 

83 """Get nodes by level in the tree. 

84 

85 Returns: 

86 List[List[Node]]: List of lists of nodes at each level 

87 """ 

88 result: list[list[Node[T]]] = [] 

89 current_level: list[Node[T]] = [self] 

90 

91 while current_level: 

92 result.append(current_level) 

93 next_level = [] 

94 

95 for node in current_level: 

96 if node.left: 

97 next_level.append(node.left) 

98 if node.right: 

99 next_level.append(node.right) 

100 

101 current_level = next_level 

102 

103 return result 

104 

105 @property 

106 def leaf_count(self) -> int: 

107 """Count the number of leaf nodes in the tree. 

108 

109 Returns: 

110 int: Number of leaf nodes 

111 """ 

112 return len(self.leaves) 

113 

114 @property 

115 def size(self) -> int: 

116 """Count the total number of nodes in the tree. 

117 

118 Returns: 

119 int: Total number of nodes 

120 """ 

121 # Counts via __iter__, which is already an iterative level-order walk, so 

122 # this inherits its freedom from the recursion limit. 

123 return sum(1 for _ in self) 

124 

125 def __iter__(self) -> Iterator[Node[T]]: 

126 """Iterate through all nodes in the tree in level-order. 

127 

128 Returns: 

129 Iterator[Node]: Iterator over all nodes 

130 """ 

131 queue: deque[Node[T]] = deque([self]) 

132 while queue: 

133 node = queue.popleft() 

134 yield node 

135 if node.left: 

136 queue.append(node.left) 

137 if node.right: 

138 queue.append(node.right)