AST Utilities

def ast_parse(contents_text: str) -> ast.Module:
    """
    Parse Python source code into AST module.
    
    Args:
        contents_text: Python source code to parse
        
    Returns:
        AST Module object
        
    Notes:
        - Suppresses warnings during parsing
        - Encodes text to bytes for ast.parse()
        - Used by core engine for plugin processing
    """

def ast_to_offset(node: ast.expr | ast.stmt) -> Offset:
    """
    Convert AST node position to tokenize offset.
    
    Args:
        node: AST expression or statement node
        
    Returns:
        Offset object with line and column information
        
    Usage:
        Used by plugins to map AST nodes to token positions
        for applying transformations at correct locations.
    """

def is_name_attr(
    node: ast.AST,
    imports: dict[str, set[str]], 
    mods: tuple[str, ...],
    names: Container[str]
) -> bool:
    """
    Check if node matches imported name or attribute pattern.
    
    Args:
        node: AST node to check
        imports: Import tracking dictionary
        mods: Module names to check
        names: Name set to match against
        
    Returns:
        True if node matches pattern
        
    Patterns matched:
        - ast.Name: Direct name usage (imported with 'from')
        - ast.Attribute: Module.name usage (imported with 'import')
    """

def has_starargs(call: ast.Call) -> bool:
    """
    Check if function call has star arguments.
    
    Args:
        call: AST Call node to analyze
        
    Returns:
        True if call has *args or **kwargs
        
    Notes:
        Used by plugins to avoid transforming calls with
        dynamic arguments that could change behavior.
    """

def is_type_check(node: ast.AST) -> bool:
    """
    Check if node is isinstance/issubclass call.
    
    Args:
        node: AST node to check
        
    Returns:
        True if node is isinstance() or issubclass() call
        
    Requirements:
        - Function name must be 'isinstance' or 'issubclass'  
        - Must have exactly 2 arguments
        - No star arguments allowed
    """

def contains_await(node: ast.AST) -> bool:
    """
    Check if AST node contains await expressions.
    
    Args:
        node: AST node to analyze
        
    Returns:
        True if any child node is an await expression
        
    Usage:
        Used to avoid transforming async generators or
        expressions that require async context.
    """

def is_async_listcomp(node: ast.ListComp) -> bool:
    """
    Check if list comprehension is async.
    
    Args:
        node: ListComp AST node to check
        
    Returns:
        True if comprehension uses async generators or await
        
    Detection criteria:
        - Any generator marked as async (gen.is_async)
        - Contains await expressions in any part
    """

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

@register(ast.Call)
def fix_collection_calls(state: State, node: ast.Call, parent: ast.AST):
    """Transform collection constructor calls."""
    
    # Check if this is a set() call
    if (isinstance(node.func, ast.Name) and 
        node.func.id == 'set'):
        
        # Get token offset for transformation
        offset = ast_to_offset(node)
        
        def transform_tokens(i: int, tokens: list[Token]) -> None:
            # Apply token transformation
            pass
            
        return [(offset, transform_tokens)]
    
    return []

@register(ast.Call)
def fix_mock_calls(state: State, node: ast.Call, parent: ast.AST):
    """Replace mock.Mock with unittest.mock.Mock."""
    
    if is_name_attr(
        node.func, 
        state.from_imports,
        ('mock',),
        {'Mock', 'patch', 'MagicMock'}
    ):
        # This is a mock call that can be transformed
        offset = ast_to_offset(node.func)
        # ... transformation logic
        return [(offset, transform_func)]
    
    return []

@register(ast.ListComp)
def fix_list_comprehensions(state: State, node: ast.ListComp, parent: ast.AST):
    """Transform list comprehensions to set comprehensions."""
    
    # Skip async comprehensions
    if is_async_listcomp(node):
        return []
        
    # Check if this can become a set comprehension
    if isinstance(parent, ast.Call) and isinstance(parent.func, ast.Name):
        if parent.func.id == 'set':
            # Transform set([x for x in y]) → {x for x in y}
            offset = ast_to_offset(parent)
            return [(offset, transform_func)]
    
    return []

@register(ast.Call) 
def fix_function_calls(state: State, node: ast.Call, parent: ast.AST):
    """Transform function calls safely."""
    
    # Skip calls with star arguments
    if has_starargs(node):
        return []
        
    # Skip type check calls (isinstance/issubclass)
    if is_type_check(node):
        return []
        
    # Safe to transform this call
    offset = ast_to_offset(node)
    return [(offset, transform_func)]