0
# Plugin System
1

2
Extensible plugin architecture for registering AST-based syntax transformations. The plugin system allows registration of transformation functions for specific AST node types and manages the visitor pattern for code analysis.
3

4
## Capabilities
5

6
### Plugin Registration
7

8
Register transformation functions for specific AST node types.
9

10
```python { .api }
11
def register(tp: type[AST_T]) -> Callable[[ASTFunc[AST_T]], ASTFunc[AST_T]]:
12
    """
13
    Decorator to register AST transformation function.
14
    
15
    Args:
16
        tp: AST node type to register for (e.g., ast.Call, ast.Name)
17
        
18
    Returns:
19
        Decorator function that registers the callback
20
        
21
    Usage:
22
        @register(ast.Call)
23
        def fix_set_literals(state, node, parent):
24
            # Return list of (offset, token_func) tuples
25
            return [...]
26
    """
27
```
28

29
### AST Visitor
30

31
Visit AST nodes and collect transformation callbacks.
32

33
```python { .api }
34
def visit(
35
    funcs: ASTCallbackMapping, 
36
    tree: ast.Module, 
37
    settings: Settings
38
) -> dict[Offset, list[TokenFunc]]:
39
    """
40
    Visit AST nodes and collect transformation callbacks.
41
    
42
    Args:
43
        funcs: Mapping of AST types to transformation functions
44
        tree: Parsed AST module to visit
45
        settings: Configuration settings
46
        
47
    Returns:
48
        Dictionary mapping token offsets to transformation functions
49
        
50
    Notes:
51
        - Tracks import statements for context
52
        - Manages annotation context for type hints
53
        - Processes nodes in depth-first order
54
    """
55
```
56

57
### Plugin State Management
58

59
State object passed to plugin functions during AST traversal.
60

61
```python { .api }
62
class State(NamedTuple):
63
    """
64
    Current state during AST traversal.
65
    
66
    Attributes:
67
        settings: Configuration settings for transformations
68
        from_imports: Tracked import statements by module
69
        in_annotation: Whether currently inside type annotation
70
    """
71
    settings: Settings
72
    from_imports: dict[str, set[str]]
73
    in_annotation: bool = False
74
```
75

76
### Plugin Function Registry
77

78
Global registry of plugin transformation functions.
79

80
```python { .api }
81
FUNCS: ASTCallbackMapping
82
"""
83
Global registry mapping AST node types to transformation functions.
84
Automatically populated by plugin modules using @register decorator.
85
"""
86

87
RECORD_FROM_IMPORTS: frozenset[str]
88
"""
89
Module names to track for import analysis:
90
- __future__, asyncio, collections, collections.abc
91
- functools, mmap, os, select, six, six.moves  
92
- socket, subprocess, sys, typing, typing_extensions
93
"""
94
```
95

96
## Plugin Development
97

98
### Plugin Function Signature
99

100
```python { .api }
101
ASTFunc = Callable[[State, AST_T, ast.AST], Iterable[tuple[Offset, TokenFunc]]]
102
"""
103
Plugin function signature.
104

105
Args:
106
    state: Current traversal state with settings and imports
107
    node: The AST node being visited (of registered type)
108
    parent: Parent AST node for context
109
    
110
Returns:
111
    Iterable of (offset, token_function) pairs for transformations
112
"""
113

114
TokenFunc = Callable[[int, list[Token]], None]
115
"""
116
Token transformation function signature.
117

118
Args:
119
    i: Token index in the token list
120
    tokens: Complete token list to modify in-place
121
"""
122
```
123

124
### Writing a Plugin
125

126
```python
127
from pyupgrade._data import register, State
128
from pyupgrade._ast_helpers import ast_to_offset
129

130
@register(ast.Call)
131
def fix_set_literals(state: State, node: ast.Call, parent: ast.AST):
132
    """Convert set([...]) to {...}."""
133
    # Check if this is a set() call
134
    if (isinstance(node.func, ast.Name) and 
135
        node.func.id == 'set' and
136
        len(node.args) == 1):
137
        
138
        offset = ast_to_offset(node)
139
        
140
        def token_callback(i: int, tokens: list[Token]) -> None:
141
            # Find and replace the set([...]) pattern
142
            # Implementation details...
143
            pass
144
            
145
        return [(offset, token_callback)]
146
    
147
    return []
148
```
149

150
### Plugin Auto-Loading
151

152
```python { .api }
153
def _import_plugins() -> None:
154
    """
155
    Automatically discover and import all plugin modules.
156
    
157
    Walks the _plugins package and imports all modules,
158
    which triggers their @register decorators to populate FUNCS.
159
    """
160
```
161

162
## Built-in Plugin Categories
163

164
### Collection Transformations
165
- **set_literals**: `set([1, 2])` → `{1, 2}`
166
- **dict_literals**: `dict([(a, b)])` → `{a: b}`
167
- **native_literals**: `list()` → `[]`
168

169
### String and Format Transformations  
170
- **fstrings**: `"{}".format(x)` → `f"{x}"`
171
- **percent_format**: `"%s" % x` → `"{}".format(x)`
172

173
### Type Annotation Upgrades
174
- **typing_pep585**: `List[int]` → `list[int]` (Python 3.9+)
175
- **typing_pep604**: `Union[int, str]` → `int | str` (Python 3.10+)
176

177
### Legacy Compatibility Removal
178
- **six_simple**: Remove six compatibility code  
179
- **mock**: `mock.Mock` → `unittest.mock.Mock`
180
- **unittest_aliases**: Update deprecated unittest methods
181

182
### Code Modernization
183
- **subprocess_run**: Update subprocess.call patterns
184
- **datetime_utc_alias**: `datetime.timezone.utc` simplification
185
- **collections_abc**: Move imports from collections to collections.abc
186

187
## Import Tracking
188

189
The plugin system tracks import statements to make context-aware transformations:
190

191
```python
192
# Tracked imports enable smart transformations
193
from typing import List, Dict
194
from collections import defaultdict
195

196
# Plugin can detect these imports and transform:
197
# List[int] → list[int] (if min_version >= (3, 9))
198
# Dict[str, int] → dict[str, int] 
199
```
200

201
### Import Context Usage
202

203
```python
204
@register(ast.Subscript)  
205
def fix_typing_generics(state: State, node: ast.Subscript, parent: ast.AST):
206
    """Replace typing generics with builtin equivalents."""
207
    if (state.settings.min_version >= (3, 9) and
208
        isinstance(node.value, ast.Name) and
209
        node.value.id in state.from_imports.get('typing', set())):
210
        
211
        # Safe to transform List[T] → list[T]
212
        # ... transformation logic
213
        pass
214
```
215

216
## Plugin Execution Flow
217

218
1. **Registration Phase**: Plugins use `@register` to populate `FUNCS`
219
2. **AST Parsing**: Source code parsed into AST tree
220
3. **Visitor Phase**: `visit()` traverses AST, calling registered plugins
221
4. **Collection Phase**: Plugins return `(offset, token_func)` pairs  
222
5. **Application Phase**: Token functions applied in reverse offset order
223
6. **Code Generation**: Modified tokens converted back to source code