0
# AST Utilities
1
2
Helper functions for working with Python AST nodes during transformations. These utilities provide common operations needed by plugins and the core transformation engine.
3
4
## Capabilities
5
6
### AST Parsing
7
8
Parse Python source code into AST with warning suppression.
9
10
```python { .api }
11
def ast_parse(contents_text: str) -> ast.Module:
12
"""
13
Parse Python source code into AST module.
14
15
Args:
16
contents_text: Python source code to parse
17
18
Returns:
19
AST Module object
20
21
Notes:
22
- Suppresses warnings during parsing
23
- Encodes text to bytes for ast.parse()
24
- Used by core engine for plugin processing
25
"""
26
```
27
28
### Position Conversion
29
30
Convert AST node positions to tokenize offsets.
31
32
```python { .api }
33
def ast_to_offset(node: ast.expr | ast.stmt) -> Offset:
34
"""
35
Convert AST node position to tokenize offset.
36
37
Args:
38
node: AST expression or statement node
39
40
Returns:
41
Offset object with line and column information
42
43
Usage:
44
Used by plugins to map AST nodes to token positions
45
for applying transformations at correct locations.
46
"""
47
```
48
49
### Import and Attribute Matching
50
51
Check if AST node matches imported names or attributes.
52
53
```python { .api }
54
def is_name_attr(
55
node: ast.AST,
56
imports: dict[str, set[str]],
57
mods: tuple[str, ...],
58
names: Container[str]
59
) -> bool:
60
"""
61
Check if node matches imported name or attribute pattern.
62
63
Args:
64
node: AST node to check
65
imports: Import tracking dictionary
66
mods: Module names to check
67
names: Name set to match against
68
69
Returns:
70
True if node matches pattern
71
72
Patterns matched:
73
- ast.Name: Direct name usage (imported with 'from')
74
- ast.Attribute: Module.name usage (imported with 'import')
75
"""
76
```
77
78
## Function Call Analysis
79
80
### Star Arguments Detection
81
82
Check if function call has star or keyword arguments.
83
84
```python { .api }
85
def has_starargs(call: ast.Call) -> bool:
86
"""
87
Check if function call has star arguments.
88
89
Args:
90
call: AST Call node to analyze
91
92
Returns:
93
True if call has *args or **kwargs
94
95
Notes:
96
Used by plugins to avoid transforming calls with
97
dynamic arguments that could change behavior.
98
"""
99
```
100
101
### Type Check Detection
102
103
Identify isinstance/issubclass calls.
104
105
```python { .api }
106
def is_type_check(node: ast.AST) -> bool:
107
"""
108
Check if node is isinstance/issubclass call.
109
110
Args:
111
node: AST node to check
112
113
Returns:
114
True if node is isinstance() or issubclass() call
115
116
Requirements:
117
- Function name must be 'isinstance' or 'issubclass'
118
- Must have exactly 2 arguments
119
- No star arguments allowed
120
"""
121
```
122
123
## Async Code Analysis
124
125
### Await Expression Detection
126
127
Check if AST subtree contains await expressions.
128
129
```python { .api }
130
def contains_await(node: ast.AST) -> bool:
131
"""
132
Check if AST node contains await expressions.
133
134
Args:
135
node: AST node to analyze
136
137
Returns:
138
True if any child node is an await expression
139
140
Usage:
141
Used to avoid transforming async generators or
142
expressions that require async context.
143
"""
144
```
145
146
### Async List Comprehension Detection
147
148
Identify async list comprehensions.
149
150
```python { .api }
151
def is_async_listcomp(node: ast.ListComp) -> bool:
152
"""
153
Check if list comprehension is async.
154
155
Args:
156
node: ListComp AST node to check
157
158
Returns:
159
True if comprehension uses async generators or await
160
161
Detection criteria:
162
- Any generator marked as async (gen.is_async)
163
- Contains await expressions in any part
164
"""
165
```
166
167
## Usage Examples
168
169
### Plugin Integration
170
171
```python
172
from pyupgrade._ast_helpers import ast_to_offset, is_name_attr
173
from pyupgrade._data import register, State
174
175
@register(ast.Call)
176
def fix_collection_calls(state: State, node: ast.Call, parent: ast.AST):
177
"""Transform collection constructor calls."""
178
179
# Check if this is a set() call
180
if (isinstance(node.func, ast.Name) and
181
node.func.id == 'set'):
182
183
# Get token offset for transformation
184
offset = ast_to_offset(node)
185
186
def transform_tokens(i: int, tokens: list[Token]) -> None:
187
# Apply token transformation
188
pass
189
190
return [(offset, transform_tokens)]
191
192
return []
193
```
194
195
### Import-Aware Transformations
196
197
```python
198
@register(ast.Call)
199
def fix_mock_calls(state: State, node: ast.Call, parent: ast.AST):
200
"""Replace mock.Mock with unittest.mock.Mock."""
201
202
if is_name_attr(
203
node.func,
204
state.from_imports,
205
('mock',),
206
{'Mock', 'patch', 'MagicMock'}
207
):
208
# This is a mock call that can be transformed
209
offset = ast_to_offset(node.func)
210
# ... transformation logic
211
return [(offset, transform_func)]
212
213
return []
214
```
215
216
### Async-Safe Transformations
217
218
```python
219
@register(ast.ListComp)
220
def fix_list_comprehensions(state: State, node: ast.ListComp, parent: ast.AST):
221
"""Transform list comprehensions to set comprehensions."""
222
223
# Skip async comprehensions
224
if is_async_listcomp(node):
225
return []
226
227
# Check if this can become a set comprehension
228
if isinstance(parent, ast.Call) and isinstance(parent.func, ast.Name):
229
if parent.func.id == 'set':
230
# Transform set([x for x in y]) → {x for x in y}
231
offset = ast_to_offset(parent)
232
return [(offset, transform_func)]
233
234
return []
235
```
236
237
### Function Call Analysis
238
239
```python
240
@register(ast.Call)
241
def fix_function_calls(state: State, node: ast.Call, parent: ast.AST):
242
"""Transform function calls safely."""
243
244
# Skip calls with star arguments
245
if has_starargs(node):
246
return []
247
248
# Skip type check calls (isinstance/issubclass)
249
if is_type_check(node):
250
return []
251
252
# Safe to transform this call
253
offset = ast_to_offset(node)
254
return [(offset, transform_func)]
255
```