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 

11from typing import TypeVar 

12 

13# Type for node values 

14NodeValue = int | float | str 

15T = TypeVar("T", bound="Node") 

16 

17 

18class Node: 

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

20 

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

22 that is used in the pyhrp package. 

23 

24 Attributes: 

25 value: The value of the node 

26 left: The left child node 

27 right: The right child node 

28 """ 

29 

30 def __init__(self, value: NodeValue, left: Node | None = None, right: Node | None = None): 

31 """Initialize a new Node. 

32 

33 Args: 

34 value: The value of the node 

35 left: The left child node 

36 right: The right child node 

37 """ 

38 self.value = value 

39 self.left = left 

40 self.right = right 

41 

42 @property 

43 def is_leaf(self) -> bool: 

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

45 

46 Returns: 

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

48 """ 

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

50 

51 @property 

52 def leaves(self) -> list[Node]: 

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

54 

55 Returns: 

56 List[Node]: List of all leaf nodes 

57 """ 

58 if self.is_leaf: 

59 return [self] 

60 

61 result = [] 

62 if self.left: 

63 result.extend(self.left.leaves) 

64 if self.right: 

65 result.extend(self.right.leaves) 

66 

67 return result 

68 

69 @property 

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

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

72 

73 Returns: 

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

75 """ 

76 result = [] 

77 current_level = [self] 

78 

79 while current_level: 

80 result.append(current_level) 

81 next_level = [] 

82 

83 for node in current_level: 

84 if node.left: 

85 next_level.append(node.left) 

86 if node.right: 

87 next_level.append(node.right) 

88 

89 current_level = next_level 

90 

91 return result 

92 

93 @property 

94 def leaf_count(self) -> int: 

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

96 

97 Returns: 

98 int: Number of leaf nodes 

99 """ 

100 return len(self.leaves) 

101 

102 @property 

103 def size(self) -> int: 

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

105 

106 Returns: 

107 int: Total number of nodes 

108 """ 

109 size = 1 # Count this node 

110 if self.left: 

111 size += self.left.size 

112 if self.right: 

113 size += self.right.size 

114 return size 

115 

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

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

118 

119 Returns: 

120 Iterator[Node]: Iterator over all nodes 

121 """ 

122 queue: deque[Node] = deque([self]) 

123 while queue: 

124 node = queue.popleft() 

125 yield node 

126 if node.left: 

127 queue.append(node.left) 

128 if node.right: 

129 queue.append(node.right)