0
# Token Manipulation
1
2
Comprehensive token-level manipulation utilities for precise code transformations. These utilities provide fine-grained control over token streams for implementing complex syntax changes.
3
4
## Constants
5
6
Core constants used for token classification and manipulation.
7
8
```python { .api }
9
_OPENING: frozenset[str]
10
"""Opening bracket characters: '([{'"""
11
12
_CLOSING: frozenset[str]
13
"""Closing bracket characters: ')]}' """
14
15
KEYWORDS: frozenset[str]
16
"""Python keyword set from keyword.kwlist"""
17
```
18
19
## Capabilities
20
21
### Token Type Detection
22
23
Basic token classification functions.
24
25
```python { .api }
26
def is_open(token: Token) -> bool:
27
"""
28
Check if token is opening bracket/parenthesis.
29
30
Args:
31
token: Token to check
32
33
Returns:
34
True if token is '(', '[', or '{'
35
"""
36
37
def is_close(token: Token) -> bool:
38
"""
39
Check if token is closing bracket/parenthesis.
40
41
Args:
42
token: Token to check
43
44
Returns:
45
True if token is ')', ']', or '}'
46
"""
47
48
def immediately_paren(func: str, tokens: list[Token], i: int) -> bool:
49
"""
50
Check if function name is immediately followed by parenthesis.
51
52
Args:
53
func: Function name to match
54
tokens: Token list
55
i: Index to check
56
57
Returns:
58
True if tokens[i] == func and tokens[i+1] == '('
59
"""
60
```
61
62
### Token Search Functions
63
64
Find specific tokens in token streams.
65
66
```python { .api }
67
def find_name(tokens: list[Token], i: int, src: str) -> int:
68
"""
69
Find next NAME token with specific source.
70
71
Args:
72
tokens: Token list to search
73
i: Starting index
74
src: Source string to match
75
76
Returns:
77
Index of matching token
78
"""
79
80
def find_op(tokens: list[Token], i: int, src: str) -> int:
81
"""
82
Find next OP token with specific source.
83
84
Args:
85
tokens: Token list to search
86
i: Starting index
87
src: Operator string to match
88
89
Returns:
90
Index of matching token
91
"""
92
93
def find_end(tokens: list[Token], i: int) -> int:
94
"""
95
Find end of statement (NEWLINE token).
96
97
Args:
98
tokens: Token list to search
99
i: Starting index
100
101
Returns:
102
Index after the NEWLINE token
103
"""
104
105
def find_call(tokens: list[Token], i: int) -> int:
106
"""
107
Find function call opening parenthesis.
108
109
Args:
110
tokens: Token list to search
111
i: Starting index
112
113
Returns:
114
Index of opening parenthesis for function call
115
116
Notes:
117
Handles nested expressions like ("something").method(...)
118
"""
119
```
120
121
### Bracket and Block Processing
122
123
Handle bracket matching and code blocks.
124
125
```python { .api }
126
def find_closing_bracket(tokens: list[Token], i: int) -> int:
127
"""
128
Find matching closing bracket for opening bracket.
129
130
Args:
131
tokens: Token list
132
i: Index of opening bracket
133
134
Returns:
135
Index of matching closing bracket
136
137
Notes:
138
Handles nested brackets correctly
139
"""
140
141
def find_block_start(tokens: list[Token], i: int) -> int:
142
"""
143
Find colon starting code block.
144
145
Args:
146
tokens: Token list
147
i: Starting search index
148
149
Returns:
150
Index of colon token starting block
151
"""
152
153
class Block(NamedTuple):
154
"""
155
Code block boundaries in token stream.
156
157
Attributes:
158
start: Block start index
159
colon: Colon token index
160
block: Block content start index
161
end: Block end index
162
line: True if single-line block
163
"""
164
start: int
165
colon: int
166
block: int
167
end: int
168
line: bool
169
170
@classmethod
171
def find(cls, tokens: list[Token], i: int, trim_end: bool = False) -> 'Block':
172
"""Find code block starting at index i."""
173
174
def dedent(self, tokens: list[Token]) -> None:
175
"""Dedent block content by removing common indentation."""
176
177
def replace_condition(self, tokens: list[Token], new: list[Token]) -> None:
178
"""Replace block condition with new tokens."""
179
```
180
181
## Function Call Processing
182
183
### Argument Parsing
184
185
Parse function call arguments from token stream.
186
187
```python { .api }
188
def parse_call_args(tokens: list[Token], i: int) -> tuple[list[tuple[int, int]], int]:
189
"""
190
Parse function call arguments from tokens.
191
192
Args:
193
tokens: Token list
194
i: Index of opening parenthesis
195
196
Returns:
197
Tuple of (argument_ranges, end_index)
198
- argument_ranges: List of (start, end) indices for each argument
199
- end_index: Index after closing parenthesis
200
"""
201
202
def arg_str(tokens: list[Token], start: int, end: int) -> str:
203
"""
204
Get argument string from token range.
205
206
Args:
207
tokens: Token list
208
start: Start index
209
end: End index
210
211
Returns:
212
String representation of tokens in range
213
"""
214
```
215
216
### Call Transformation
217
218
Replace function calls with templates.
219
220
```python { .api }
221
def replace_call(
222
tokens: list[Token],
223
start: int,
224
end: int,
225
args: list[tuple[int, int]],
226
tmpl: str,
227
*,
228
parens: Sequence[int] = ()
229
) -> None:
230
"""
231
Replace function call with template.
232
233
Args:
234
tokens: Token list to modify in-place
235
start: Call start index
236
end: Call end index
237
args: Argument ranges from parse_call_args
238
tmpl: Template string with {args[0]}, {args[1]}, {rest} placeholders
239
parens: Argument indices to wrap in parentheses
240
241
Notes:
242
- Handles multiline arguments safely
243
- Preserves comments and whitespace where possible
244
- Adds parentheses around arguments that contain newlines
245
"""
246
247
def find_and_replace_call(
248
i: int,
249
tokens: list[Token],
250
*,
251
template: str,
252
parens: tuple[int, ...] = ()
253
) -> None:
254
"""
255
Find function call and replace with template.
256
257
Args:
258
i: Starting token index
259
tokens: Token list to modify
260
template: Replacement template
261
parens: Argument indices to parenthesize
262
"""
263
```
264
265
## Token Modification Utilities
266
267
### Token Replacement
268
269
Replace individual tokens and arguments.
270
271
```python { .api }
272
def replace_name(i: int, tokens: list[Token], *, name: str, new: str) -> None:
273
"""
274
Replace name token with new name.
275
276
Args:
277
i: Starting token index
278
tokens: Token list to modify
279
name: Name to find and replace
280
new: Replacement name
281
"""
282
283
def delete_argument(
284
i: int,
285
tokens: list[Token],
286
func_args: Sequence[tuple[int, int]]
287
) -> None:
288
"""
289
Delete function argument from call.
290
291
Args:
292
i: Argument index to delete
293
tokens: Token list to modify
294
func_args: Argument ranges from parse_call_args
295
"""
296
297
def replace_argument(
298
i: int,
299
tokens: list[Token],
300
func_args: Sequence[tuple[int, int]],
301
*,
302
new: str
303
) -> None:
304
"""
305
Replace function argument with new content.
306
307
Args:
308
i: Argument index to replace
309
tokens: Token list to modify
310
func_args: Argument ranges
311
new: Replacement content
312
"""
313
```
314
315
### Structural Modifications
316
317
Remove code structures like braces, decorators, and base classes.
318
319
```python { .api }
320
def remove_brace(tokens: list[Token], i: int) -> None:
321
"""
322
Remove brace token and surrounding whitespace.
323
324
Args:
325
tokens: Token list to modify
326
i: Index of brace token
327
328
Notes:
329
Removes extra whitespace if brace is on its own line
330
"""
331
332
def remove_decorator(i: int, tokens: list[Token]) -> None:
333
"""
334
Remove decorator from function/class.
335
336
Args:
337
i: Index within decorator
338
tokens: Token list to modify
339
340
Notes:
341
Finds @ symbol and removes entire decorator line
342
"""
343
344
def remove_base_class(i: int, tokens: list[Token]) -> None:
345
"""
346
Remove base class from class definition.
347
348
Args:
349
i: Index within base class reference
350
tokens: Token list to modify
351
352
Notes:
353
Handles single/multiple base classes and parentheses correctly
354
"""
355
```
356
357
## Comprehension Processing
358
359
### Victim Token Identification
360
361
Identify tokens to remove for comprehension transformations.
362
363
```python { .api }
364
class Victims(NamedTuple):
365
"""
366
Token indices to remove for comprehension transformations.
367
368
Attributes:
369
starts: Opening bracket indices to remove
370
ends: Closing bracket indices to remove
371
first_comma_index: First comma index (if any)
372
arg_index: Argument start index
373
"""
374
starts: list[int]
375
ends: list[int]
376
first_comma_index: int | None
377
arg_index: int
378
379
def victims(
380
tokens: list[Token],
381
start: int,
382
arg: ast.expr,
383
gen: bool
384
) -> Victims:
385
"""
386
Find tokens to remove for comprehension transformation.
387
388
Args:
389
tokens: Token list
390
start: Starting token index
391
arg: AST expression for argument
392
gen: True if generator expression
393
394
Returns:
395
Victims object with token indices to remove
396
397
Usage:
398
Used to transform set([x for x in y]) → {x for x in y}
399
"""
400
```
401
402
## Advanced Utilities
403
404
### Constant Folding
405
406
Fold constant expressions in tuples.
407
408
```python { .api }
409
def constant_fold_tuple(i: int, tokens: list[Token]) -> None:
410
"""
411
Fold constant tuple expressions.
412
413
Args:
414
i: Starting token index
415
tokens: Token list to modify
416
417
Example:
418
isinstance(x, (int, int, str)) → isinstance(x, (int, str))
419
"""
420
```
421
422
### Indentation and Spacing
423
424
Handle indentation and whitespace.
425
426
```python { .api }
427
def has_space_before(i: int, tokens: list[Token]) -> bool:
428
"""
429
Check if token has whitespace before it.
430
431
Args:
432
i: Token index
433
tokens: Token list
434
435
Returns:
436
True if preceded by whitespace or indent
437
"""
438
439
def indented_amount(i: int, tokens: list[Token]) -> str:
440
"""
441
Get indentation string at token position.
442
443
Args:
444
i: Token index
445
tokens: Token list
446
447
Returns:
448
Indentation string (spaces/tabs)
449
450
Raises:
451
ValueError: If not at beginning of line
452
"""
453
```
454
455
## Usage Examples
456
457
### Function Call Transformation
458
459
```python
460
from pyupgrade._token_helpers import find_and_replace_call
461
462
def transform_print_calls(i: int, tokens: list[Token]) -> None:
463
"""Transform print statements to print functions."""
464
465
# Replace print x, y with print(x, y)
466
find_and_replace_call(
467
i, tokens,
468
template="print({args[0]}, {args[1]})",
469
parens=()
470
)
471
```
472
473
### Comprehension Transformation
474
475
```python
476
from pyupgrade._token_helpers import victims, remove_brace
477
478
def transform_set_comprehension(tokens: list[Token], start: int, arg: ast.expr) -> None:
479
"""Transform set([x for x in y]) to {x for x in y}."""
480
481
# Find tokens to remove
482
v = victims(tokens, start, arg, gen=False)
483
484
# Remove brackets in reverse order
485
for end_idx in reversed(v.ends):
486
remove_brace(tokens, end_idx)
487
488
for start_idx in reversed(v.starts):
489
remove_brace(tokens, start_idx)
490
```