0
# AST Node Utilities
1

2
Helper functions for working with AST nodes, including type checking, tree traversal, and node classification utilities. These functions provide support for both standard Python `ast` module and `astroid` library nodes.
3

4
## Capabilities
5

6
### Tree Traversal
7

8
Functions for walking and visiting AST trees with fine-grained control over traversal order and node inclusion.
9

10
```python { .api }
11
# Available through asttokens.util module
12
from asttokens.util import (
13
    walk, visit_tree, iter_children_ast, iter_children_astroid, 
14
    iter_children_func
15
)
16

17
def walk(node, include_joined_str=False) -> Iterator:
18
    """
19
    Pre-order traversal of AST tree yielding all descendant nodes.
20
    
21
    Parameters:
22
    - node (ast.AST): Root node to start traversal from
23
    - include_joined_str (bool): Include f-string nodes that may have broken positions
24
    
25
    Yields:
26
    ast.AST: Each node in pre-order traversal
27
    """
28

29
def visit_tree(node, previsit, postvisit):
30
    """
31
    Depth-first tree traversal with explicit stack and pre/post visit callbacks.
32
    
33
    Parameters:
34
    - node (ast.AST): Root node to traverse
35
    - previsit (callable): Function called before visiting children
36
    - postvisit (callable): Function called after visiting children
37
    """
38

39
def iter_children_ast(node, include_joined_str=False) -> Iterator:
40
    """
41
    Iterate over direct children of an ast module node.
42
    
43
    Parameters:
44
    - node (ast.AST): Node to get children from
45
    - include_joined_str (bool): Include f-string nodes
46
    
47
    Yields:
48
    ast.AST: Each child node
49
    """
50

51
def iter_children_astroid(node, include_joined_str=False) -> Iterator:
52
    """
53
    Iterate over direct children of an astroid node.
54
    
55
    Parameters:
56
    - node: Astroid node to get children from
57
    - include_joined_str (bool): Include f-string nodes
58
    
59
    Yields:
60
    Astroid node: Each child node
61
    """
62

63
def iter_children_func(node) -> callable:
64
    """
65
    Return appropriate child iterator for AST or Astroid nodes.
66
    
67
    Parameters:
68
    - node: AST or Astroid node
69
    
70
    Returns:
71
    callable: iter_children_ast or iter_children_astroid function
72
    """
73
```
74

75
#### Usage Example
76

77
```python
78
import asttokens
79
import ast
80

81
source = """
82
def greet(name):
83
    if name:
84
        return f"Hello, {name}!"
85
    return "Hello, World!"
86
"""
87

88
tree = ast.parse(source)
89

90
# Walk all nodes in tree
91
for node in asttokens.util.walk(tree):
92
    print(f"{type(node).__name__}: {getattr(node, 'lineno', 'N/A')}")
93

94
# Custom tree traversal with callbacks
95
def pre_visit(node, parent, field):
96
    print(f"Entering: {type(node).__name__}")
97

98
def post_visit(node, parent, field):
99
    print(f"Exiting: {type(node).__name__}")
100

101
asttokens.util.visit_tree(tree, pre_visit, post_visit)
102

103
# Iterate children of function node
104
func_node = tree.body[0]
105
for child in asttokens.util.iter_children_ast(func_node):
106
    print(f"Child: {type(child).__name__}")
107
```
108

109
### Node Type Checking
110

111
Functions for identifying different categories of AST nodes by their semantic role.
112

113
```python { .api }
114
# Available through asttokens.util module
115
from asttokens.util import (
116
    is_expr, is_stmt, is_module, is_joined_str, is_expr_stmt,
117
    is_constant, is_ellipsis, is_starred, is_slice, is_empty_astroid_slice
118
)
119

120
def is_expr(node) -> bool:
121
    """
122
    Check if node is an expression.
123
    
124
    Parameters:
125
    - node (ast.AST): Node to check
126
    
127
    Returns:
128
    bool: True if node is an expression
129
    """
130

131
def is_stmt(node) -> bool:
132
    """
133
    Check if node is a statement.
134
    
135
    Parameters:
136
    - node (ast.AST): Node to check
137
    
138
    Returns:
139
    bool: True if node is a statement
140
    """
141

142
def is_module(node) -> bool:
143
    """
144
    Check if node is a module.
145
    
146
    Parameters:
147
    - node (ast.AST): Node to check
148
    
149
    Returns:
150
    bool: True if node is a Module
151
    """
152

153
def is_joined_str(node) -> bool:
154
    """
155
    Check if node is an f-string (JoinedStr).
156
    
157
    Parameters:
158
    - node (ast.AST): Node to check
159
    
160
    Returns:
161
    bool: True if node is a JoinedStr (f-string)
162
    """
163

164
def is_expr_stmt(node) -> bool:
165
    """
166
    Check if node is an expression statement.
167
    
168
    Parameters:
169
    - node (ast.AST): Node to check
170
    
171
    Returns:
172
    bool: True if node is an Expr statement
173
    """
174

175
def is_constant(node) -> bool:
176
    """
177
    Check if node represents a constant value.
178
    
179
    Parameters:
180
    - node (ast.AST): Node to check
181
    
182
    Returns:
183
    bool: True if node is a constant (ast.Constant or astroid.Const)
184
    """
185

186
def is_ellipsis(node) -> bool:
187
    """
188
    Check if node represents ellipsis (...).
189
    
190
    Parameters:
191
    - node (ast.AST): Node to check
192
    
193
    Returns:
194
    bool: True if node represents ellipsis
195
    """
196

197
def is_starred(node) -> bool:
198
    """
199
    Check if node is a starred expression (*args).
200
    
201
    Parameters:
202
    - node (ast.AST): Node to check
203
    
204
    Returns:
205
    bool: True if node is a Starred expression
206
    """
207

208
def is_slice(node) -> bool:
209
    """
210
    Check if node represents a slice operation.
211
    
212
    Parameters:
213
    - node (ast.AST): Node to check
214
    
215
    Returns:
216
    bool: True if node represents slice
217
    """
218

219
def is_empty_astroid_slice(node) -> bool:
220
    """
221
    Check if node is an empty Astroid slice.
222
    
223
    Parameters:
224
    - node: Node to check (Astroid)
225
    
226
    Returns:
227
    bool: True if node is empty Astroid slice
228
    """
229
```
230

231
#### Usage Example
232

233
```python
234
import asttokens
235
import ast
236

237
source = """
238
x = 42
239
print(x)
240
y = [1, 2, 3]
241
z = y[1:3]
242
def func(*args): pass
243
"""
244

245
tree = ast.parse(source)
246

247
# Check node types
248
for node in ast.walk(tree):
249
    if asttokens.util.is_stmt(node):
250
        print(f"Statement: {type(node).__name__}")
251
    elif asttokens.util.is_expr(node):
252
        print(f"Expression: {type(node).__name__}")
253
    
254
    if asttokens.util.is_constant(node):
255
        print(f"Constant: {node.value}")
256
    
257
    if asttokens.util.is_starred(node):
258
        print("Found starred expression")
259
        
260
    if asttokens.util.is_slice(node):
261
        print("Found slice operation")
262

263
# Check specific patterns
264
assign_node = tree.body[0]  # x = 42
265
print(f"Is statement: {asttokens.util.is_stmt(assign_node)}")
266

267
constant_node = assign_node.value  # 42
268
print(f"Is constant: {asttokens.util.is_constant(constant_node)}")
269
```
270

271
### Statement and Expression Utilities
272

273
Functions for working with specific types of statements and expressions.
274

275
```python { .api }
276
# Available through asttokens.util module
277
from asttokens.util import (
278
    last_stmt, replace, annotate_fstring_nodes, fstring_positions_work
279
)
280

281
def last_stmt(node) -> ast.AST:
282
    """
283
    Get the last statement in a multi-statement node.
284
    
285
    Parameters:
286
    - node (ast.AST): Node that may contain multiple statements
287
    
288
    Returns:
289
    ast.AST: Last statement node
290
    """
291

292
def replace(text, replacements) -> str:
293
    """
294
    Apply multiple text replacements to source code.
295
    
296
    Parameters:
297
    - text (str): Original text
298
    - replacements (List[Tuple[int, int, str]]): List of (start, end, replacement) tuples
299
    
300
    Returns:
301
    str: Text with replacements applied
302
    """
303

304
def annotate_fstring_nodes(tree):
305
    """
306
    Mark f-string nodes that have broken position information.
307
    
308
    Parameters:
309
    - tree (ast.AST): AST tree to annotate
310
    """
311

312
def fstring_positions_work() -> bool:
313
    """
314
    Check if f-string positions work correctly in current Python version.
315
    
316
    Returns:
317
    bool: True if f-string positions are reliable
318
    """
319
```
320

321
#### Usage Example
322

323
```python
324
import asttokens
325
import ast
326

327
source = """
328
if True:
329
    print("first")
330
    print("second")
331
    x = 42
332
"""
333

334
tree = ast.parse(source)
335
if_node = tree.body[0]
336

337
# Get last statement in if block
338
last = asttokens.util.last_stmt(if_node.body[0])
339
print(f"Last statement: {type(last).__name__}")
340

341
# Check f-string support
342
if asttokens.util.fstring_positions_work():
343
    print("F-string positions work correctly")
344
    
345
# Apply text replacements
346
original = "hello world"
347
replacements = [(0, 5, "hi"), (6, 11, "there")]
348
result = asttokens.util.replace(original, replacements)
349
print(result)  # "hi there"
350

351
# Annotate f-strings in tree (modifies tree in-place)
352
asttokens.util.annotate_fstring_nodes(tree)
353
```
354

355
### Additional Utility Functions
356

357
Additional utility functions for working with tokens and text manipulation.
358

359
```python { .api }
360
# Available through asttokens.util module
361
from asttokens.util import combine_tokens
362

363
def combine_tokens(group) -> list:
364
    """
365
    Combine consecutive NAME, NUMBER, and ERRORTOKEN tokens.
366
    
367
    Parameters:
368
    - group (List[tokenize.TokenInfo]): List of tokens to combine
369
    
370
    Returns:
371
    List[tokenize.TokenInfo]: Combined token list
372
    """
373
```
374

375
#### Usage Example
376

377
```python
378
import asttokens
379
import asttokens.util
380

381
# The combine_tokens function is used internally for token processing
382
# but is available for advanced token manipulation scenarios
383
source = "hello_world123"
384
tokens = list(asttokens.util.generate_tokens(source))
385
print([t.string for t in tokens])
386
```