tessl/pypi-pyupgrade

A tool and pre-commit hook to automatically upgrade Python syntax for newer versions of the language.

—

Pending

Overview

Eval results

Files

Plugin System

Name: tessl/pypi-pyupgrade
Author: tessl

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.

Capabilities

Plugin Registration

def register(tp: type[AST_T]) -> Callable[[ASTFunc[AST_T]], ASTFunc[AST_T]]:
    """
    Decorator to register AST transformation function.
    
    Args:
        tp: AST node type to register for (e.g., ast.Call, ast.Name)
        
    Returns:
        Decorator function that registers the callback
        
    Usage:
        @register(ast.Call)
        def fix_set_literals(state, node, parent):
            # Return list of (offset, token_func) tuples
            return [...]
    """

AST Visitor

Visit AST nodes and collect transformation callbacks.

def visit(
    funcs: ASTCallbackMapping, 
    tree: ast.Module, 
    settings: Settings
) -> dict[Offset, list[TokenFunc]]:
    """
    Visit AST nodes and collect transformation callbacks.
    
    Args:
        funcs: Mapping of AST types to transformation functions
        tree: Parsed AST module to visit
        settings: Configuration settings
        
    Returns:
        Dictionary mapping token offsets to transformation functions
        
    Notes:
        - Tracks import statements for context
        - Manages annotation context for type hints
        - Processes nodes in depth-first order
    """

Plugin State Management

State object passed to plugin functions during AST traversal.

class State(NamedTuple):
    """
    Current state during AST traversal.
    
    Attributes:
        settings: Configuration settings for transformations
        from_imports: Tracked import statements by module
        in_annotation: Whether currently inside type annotation
    """
    settings: Settings
    from_imports: dict[str, set[str]]
    in_annotation: bool = False

Plugin Function Registry

Global registry of plugin transformation functions.

FUNCS: ASTCallbackMapping
"""
Global registry mapping AST node types to transformation functions.
Automatically populated by plugin modules using @register decorator.
"""

RECORD_FROM_IMPORTS: frozenset[str]
"""
Module names to track for import analysis:
- __future__, asyncio, collections, collections.abc
- functools, mmap, os, select, six, six.moves  
- socket, subprocess, sys, typing, typing_extensions
"""

Plugin Development

Plugin Function Signature

ASTFunc = Callable[[State, AST_T, ast.AST], Iterable[tuple[Offset, TokenFunc]]]
"""
Plugin function signature.

Args:
    state: Current traversal state with settings and imports
    node: The AST node being visited (of registered type)
    parent: Parent AST node for context
    
Returns:
    Iterable of (offset, token_function) pairs for transformations
"""

TokenFunc = Callable[[int, list[Token]], None]
"""
Token transformation function signature.

Args:
    i: Token index in the token list
    tokens: Complete token list to modify in-place
"""

Writing a Plugin

from pyupgrade._data import register, State
from pyupgrade._ast_helpers import ast_to_offset

@register(ast.Call)
def fix_set_literals(state: State, node: ast.Call, parent: ast.AST):
    """Convert set([...]) to {...}."""
    # Check if this is a set() call
    if (isinstance(node.func, ast.Name) and 
        node.func.id == 'set' and
        len(node.args) == 1):
        
        offset = ast_to_offset(node)
        
        def token_callback(i: int, tokens: list[Token]) -> None:
            # Find and replace the set([...]) pattern
            # Implementation details...
            pass
            
        return [(offset, token_callback)]
    
    return []

Plugin Auto-Loading

def _import_plugins() -> None:
    """
    Automatically discover and import all plugin modules.
    
    Walks the _plugins package and imports all modules,
    which triggers their @register decorators to populate FUNCS.
    """

Built-in Plugin Categories

Collection Transformations

set_literals: set([1, 2]) → {1, 2}
dict_literals: dict([(a, b)]) → {a: b}
native_literals: list() → []

String and Format Transformations

fstrings: "{}".format(x) → f"{x}"
percent_format: "%s" % x → "{}".format(x)

Type Annotation Upgrades

typing_pep585: List[int] → list[int] (Python 3.9+)
typing_pep604: Union[int, str] → int | str (Python 3.10+)

Legacy Compatibility Removal

six_simple: Remove six compatibility code
mock: mock.Mock → unittest.mock.Mock
unittest_aliases: Update deprecated unittest methods

Code Modernization

subprocess_run: Update subprocess.call patterns
datetime_utc_alias: datetime.timezone.utc simplification
collections_abc: Move imports from collections to collections.abc

Import Tracking

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

# Tracked imports enable smart transformations
from typing import List, Dict
from collections import defaultdict

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

Import Context Usage

@register(ast.Subscript)  
def fix_typing_generics(state: State, node: ast.Subscript, parent: ast.AST):
    """Replace typing generics with builtin equivalents."""
    if (state.settings.min_version >= (3, 9) and
        isinstance(node.value, ast.Name) and
        node.value.id in state.from_imports.get('typing', set())):
        
        # Safe to transform List[T] → list[T]
        # ... transformation logic
        pass