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 TYPE_CHECKING 

7 

8from _griffe.exceptions import LastNodeError 

9 

10if TYPE_CHECKING: 

11 from collections.abc import Iterator 

12 

13 

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

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

16 

17 Parameters: 

18 node: The AST node. 

19 

20 Returns: 

21 The node kind. 

22 """ 

23 return node.__class__.__name__.lower() 

24 

25 

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

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

28 

29 Parameters: 

30 node: The AST node. 

31 

32 Yields: 

33 The node children. 

34 """ 

35 for field_name in node._fields: 

36 try: 

37 field = getattr(node, field_name) 

38 except AttributeError: 

39 continue 

40 if isinstance(field, AST): 

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

42 yield field 

43 elif isinstance(field, list): 

44 for child in field: 

45 if isinstance(child, AST): 

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

47 yield child 

48 

49 

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

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

52 

53 Parameters: 

54 node: The AST node. 

55 

56 Yields: 

57 The previous siblings. 

58 """ 

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

60 if sibling is not node: 

61 yield sibling 

62 else: 

63 return 

64 

65 

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

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

68 

69 Parameters: 

70 node: The AST node. 

71 

72 Yields: 

73 The next siblings. 

74 """ 

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

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

77 if sibling is node: 

78 break 

79 yield from siblings 

80 

81 

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

83 """Return the siblings of this node. 

84 

85 Parameters: 

86 node: The AST node. 

87 

88 Yields: 

89 The siblings. 

90 """ 

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

92 for sibling in siblings: 

93 if sibling is not node: 

94 yield sibling 

95 else: 

96 break 

97 yield from siblings 

98 

99 

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

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

102 

103 Parameters: 

104 node: The AST node. 

105 

106 Raises: 

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

108 

109 Returns: 

110 The sibling. 

111 """ 

112 try: 

113 *_, last = ast_previous_siblings(node) 

114 except ValueError: 

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

116 return last 

117 

118 

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

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

121 

122 Parameters: 

123 node: The AST node. 

124 

125 Raises: 

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

127 

128 Returns: 

129 The sibling. 

130 """ 

131 try: 

132 return next(ast_next_siblings(node)) 

133 except StopIteration: 

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

135 

136 

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

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

139 

140 Parameters: 

141 node: The AST node. 

142 

143 Raises: 

144 LastNodeError: When the node does not have children. 

145 

146 Returns: 

147 The child. 

148 """ 

149 try: 

150 return next(ast_children(node)) 

151 except StopIteration as error: 

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

153 

154 

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

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

157 

158 Parameters: 

159 node: The AST node. 

160 

161 Raises: 

162 LastNodeError: When the node does not have children. 

163 

164 Returns: 

165 The child. 

166 """ 

167 try: 

168 *_, last = ast_children(node) 

169 except ValueError as error: 

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

171 return last