Coverage for src/_griffe/agents/nodes/ast.py: 57.32%

1# This module contains utilities for navigating AST nodes. 

2 

3from __future__ import annotations 

4 

5from ast import AST 

6from typing import Iterator 

7 

8from _griffe.exceptions import LastNodeError 

9 

10 

11def ast_kind(node: AST) -> str: 

12 """Return the kind of an AST node. 

13 

14 Parameters: 

15 node: The AST node. 

16 

17 Returns: 

18 The node kind. 

19 """ 

20 return node.__class__.__name__.lower() 

21 

22 

23def ast_children(node: AST) -> Iterator[AST]: 

24 """Return the children of an AST node. 

25 

26 Parameters: 

27 node: The AST node. 

28 

29 Yields: 

30 The node children. 

31 """ 

32 for field_name in node._fields: 

33 try: 

34 field = getattr(node, field_name) 

35 except AttributeError: 

36 continue 

37 if isinstance(field, AST): 

38 field.parent = node # type: ignore[attr-defined] 

39 yield field 

40 elif isinstance(field, list): 

41 for child in field: 

42 if isinstance(child, AST): 

43 child.parent = node # type: ignore[attr-defined] 

44 yield child 

45 

46 

47def ast_previous_siblings(node: AST) -> Iterator[AST]: 

48 """Return the previous siblings of this node, starting from the closest. 

49 

50 Parameters: 

51 node: The AST node. 

52 

53 Yields: 

54 The previous siblings. 

55 """ 

56 for sibling in ast_children(node.parent): # type: ignore[attr-defined] 

57 if sibling is not node: 

58 yield sibling 

59 else: 

60 return 

61 

62 

63def ast_next_siblings(node: AST) -> Iterator[AST]: 

64 """Return the next siblings of this node, starting from the closest. 

65 

66 Parameters: 

67 node: The AST node. 

68 

69 Yields: 

70 The next siblings. 

71 """ 

72 siblings = ast_children(node.parent) # type: ignore[attr-defined] 

73 for sibling in siblings: 73 ↛ 76line 73 didn't jump to line 76 because the loop on line 73 didn't complete

74 if sibling is node: 

75 break 

76 yield from siblings 

77 

78 

79def ast_siblings(node: AST) -> Iterator[AST]: 

80 """Return the siblings of this node. 

81 

82 Parameters: 

83 node: The AST node. 

84 

85 Yields: 

86 The siblings. 

87 """ 

88 siblings = ast_children(node.parent) # type: ignore[attr-defined] 

89 for sibling in siblings: 

90 if sibling is not node: 

91 yield sibling 

92 else: 

93 break 

94 yield from siblings 

95 

96 

97def ast_previous(node: AST) -> AST: 

98 """Return the previous sibling of this node. 

99 

100 Parameters: 

101 node: The AST node. 

102 

103 Raises: 

104 LastNodeError: When the node does not have previous siblings. 

105 

106 Returns: 

107 The sibling. 

108 """ 

109 try: 

110 *_, last = ast_previous_siblings(node) 

111 except ValueError: 

112 raise LastNodeError("there is no previous node") from None 

113 return last 

114 

115 

116def ast_next(node: AST) -> AST: 

117 """Return the next sibling of this node. 

118 

119 Parameters: 

120 node: The AST node. 

121 

122 Raises: 

123 LastNodeError: When the node does not have next siblings. 

124 

125 Returns: 

126 The sibling. 

127 """ 

128 try: 

129 return next(ast_next_siblings(node)) 

130 except StopIteration: 

131 raise LastNodeError("there is no next node") from None 

132 

133 

134def ast_first_child(node: AST) -> AST: 

135 """Return the first child of this node. 

136 

137 Parameters: 

138 node: The AST node. 

139 

140 Raises: 

141 LastNodeError: When the node does not have children. 

142 

143 Returns: 

144 The child. 

145 """ 

146 try: 

147 return next(ast_children(node)) 

148 except StopIteration as error: 

149 raise LastNodeError("there are no children node") from error 

150 

151 

152def ast_last_child(node: AST) -> AST: 

153 """Return the lasts child of this node. 

154 

155 Parameters: 

156 node: The AST node. 

157 

158 Raises: 

159 LastNodeError: When the node does not have children. 

160 

161 Returns: 

162 The child. 

163 """ 

164 try: 

165 *_, last = ast_children(node) 

166 except ValueError as error: 

167 raise LastNodeError("there are no children node") from error 

168 return last