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 if self.is_leaf: 

62 return [self] 

63 

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

65 if self.left: 

66 result.extend(self.left.leaves) 

67 if self.right: 

68 result.extend(self.right.leaves) 

69 

70 return result 

71 

72 @property 

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

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

75 

76 Returns: 

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

78 """ 

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

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

81 

82 while current_level: 

83 result.append(current_level) 

84 next_level = [] 

85 

86 for node in current_level: 

87 if node.left: 

88 next_level.append(node.left) 

89 if node.right: 

90 next_level.append(node.right) 

91 

92 current_level = next_level 

93 

94 return result 

95 

96 @property 

97 def leaf_count(self) -> int: 

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

99 

100 Returns: 

101 int: Number of leaf nodes 

102 """ 

103 return len(self.leaves) 

104 

105 @property 

106 def size(self) -> int: 

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

108 

109 Returns: 

110 int: Total number of nodes 

111 """ 

112 size = 1 # Count this node 

113 if self.left: 

114 size += self.left.size 

115 if self.right: 

116 size += self.right.size 

117 return size 

118 

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

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

121 

122 Returns: 

123 Iterator[Node]: Iterator over all nodes 

124 """ 

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

126 while queue: 

127 node = queue.popleft() 

128 yield node 

129 if node.left: 

130 queue.append(node.left) 

131 if node.right: 

132 queue.append(node.right)