Python Syntax Elements

class PythonMixin:
    """
    Mixin providing Python-specific functionality to nodes and leaves.
    """

class PythonLeaf(PythonMixin, Leaf):
    """
    Base class for Python leaf nodes (tokens).
    
    Inherits all functionality from Leaf with Python-specific enhancements.
    """

class PythonBaseNode(PythonMixin, BaseNode):
    """
    Base class for Python nodes with children.
    
    Inherits all functionality from BaseNode with Python-specific enhancements.
    """

class PythonNode(PythonMixin, Node):
    """
    Standard Python node implementation.
    
    Used for most Python syntax constructs.
    """

class PythonErrorNode(PythonMixin, ErrorNode):
    """
    Error node for invalid Python syntax.
    
    Contains partial or invalid syntax with error recovery.
    """

class PythonErrorLeaf(ErrorLeaf, PythonLeaf):
    """
    Error leaf for invalid Python tokens.
    
    Contains partial or invalid tokens with error recovery.
    """

class Module(Scope):
    """
    The top-level module node representing a Python file.
    
    Attributes:
        type (str): Always 'file_input'
    """
    
    def get_used_names(self):
        """
        Get all name usages in this module.
        
        Returns:
            UsedNamesMapping: Immutable mapping from name strings to lists of Name nodes
        """
    
    def iter_imports(self):
        """
        Get all import statements in this module.
        
        Yields:
            Import: ImportName and ImportFrom nodes
        """
    
    def iter_funcdefs(self):
        """
        Get all function definitions in this module.
        
        Yields:
            Function: Function definition nodes
        """
    
    def iter_classdefs(self):
        """
        Get all class definitions in this module.
        
        Yields:
            Class: Class definition nodes
        """

class Scope(PythonBaseNode):
    """
    Base class for scoped constructs (functions, classes, modules).
    """
    
    def iter_funcdefs(self):
        """Get function definitions within this scope."""
    
    def iter_classdefs(self):
        """Get class definitions within this scope."""
    
    def iter_imports(self):
        """Get import statements within this scope."""
    
    def get_suite(self):
        """
        Get the executable body of this scope.
        
        Returns:
            BaseNode: Suite node containing the body statements
        """

import parso

code = '''
import os
import sys
from pathlib import Path

def function1():
    pass

class MyClass:
    def method(self):
        pass

def function2():
    return 42
'''

module = parso.parse(code)

# Get all imports
imports = list(module.iter_imports())
print(f"Found {len(imports)} import statements")
for imp in imports:
    if imp.type == 'import_name':
        names = [path[-1].value for path in imp.get_paths()]
        print(f"Import: {', '.join(names)}")
    elif imp.type == 'import_from':
        from_names = [n.value for n in imp.get_from_names()]
        defined = [n.value for n in imp.get_defined_names()]
        print(f"From {'.'.join(from_names)} import {', '.join(defined)}")

# Get all functions
functions = list(module.iter_funcdefs())
print(f"Found {len(functions)} functions:")
for func in functions:
    print(f"  {func.name.value}() at line {func.start_pos[0]}")

# Get all classes
classes = list(module.iter_classdefs())
print(f"Found {len(classes)} classes:")
for cls in classes:
    print(f"  {cls.name.value} at line {cls.start_pos[0]}")

# Analyze name usage
used_names = module.get_used_names()
print(f"Name 'os' used {len(used_names.get('os', []))} times")

class Function(ClassOrFunc):
    """
    Function definition node.
    
    Attributes:
        type (str): 'funcdef'
        name (Name): Function name node
    """
    
    def get_params(self):
        """
        Get function parameters.
        
        Returns:
            list[Param]: List of parameter objects
        """
    
    def iter_yield_exprs(self):
        """
        Find yield expressions in function body.
        
        Yields:
            YieldExpr: Yield expression nodes
        """
    
    def iter_return_stmts(self):
        """
        Find return statements in function body.
        
        Yields:
            ReturnStmt: Return statement nodes
        """
    
    def iter_raise_stmts(self):
        """
        Find raise statements in function body.
        
        Yields:
            KeywordStatement: Raise statement nodes
        """
    
    def is_generator(self):
        """
        Check if function is a generator.
        
        Returns:
            bool: True if function contains yield expressions
        """
    
    @property
    def annotation(self):
        """
        Get return type annotation.
        
        Returns:
            NodeOrLeaf | None: Return annotation or None
        """

class Class(ClassOrFunc):
    """
    Class definition node.
    
    Attributes:
        type (str): 'classdef'
        name (Name): Class name node
    """
    
    def get_super_arglist(self):
        """
        Get superclass argument list.
        
        Returns:
            BaseNode | None: Arglist node with superclasses or None
        """

class Lambda(Function):
    """
    Lambda expression node.
    
    Attributes:
        type (str): 'lambdef'
    """
    
    @property
    def name(self):
        """
        Lambda name property.
        
        Raises:
            AttributeError: Lambdas don't have names
        """
    
    @property
    def annotation(self):
        """
        Lambda annotation property.
        
        Returns:
            None: Lambdas don't have return annotations
        """

import parso

code = '''
def regular_func(a: int, b: str = "default") -> str:
    """A regular function."""
    return f"{a}: {b}"

def generator_func():
    yield 1
    yield 2
    return "done"

class Parent:
    pass

class Child(Parent, dict):
    def method(self, x, y=None):
        if x:
            return x
        raise ValueError("Invalid input")

lambda_expr = lambda x, y: x + y
'''

module = parso.parse(code)

# Analyze function definitions
for func in module.iter_funcdefs():
    print(f"\nFunction: {func.name.value}")
    
    # Check if it's a generator
    if func.is_generator():
        print("  Type: Generator")
        yields = list(func.iter_yield_exprs())
        print(f"  Yield expressions: {len(yields)}")
    else:
        print("  Type: Regular function")
    
    # Get parameters
    params = func.get_params()
    print(f"  Parameters: {len(params)}")
    for param in params:
        param_info = [param.name.value]
        if param.annotation:
            param_info.append(f"annotation={param.annotation.get_code()}")
        if param.default:
            param_info.append(f"default={param.default.get_code()}")
        if param.star_count:
            param_info.append(f"star_count={param.star_count}")
        print(f"    {', '.join(param_info)}")
    
    # Check return annotation
    if func.annotation:
        print(f"  Return annotation: {func.annotation.get_code()}")
    
    # Count return statements
    returns = list(func.iter_return_stmts())
    print(f"  Return statements: {len(returns)}")

# Analyze class definitions
for cls in module.iter_classdefs():
    print(f"\nClass: {cls.name.value}")
    
    # Get superclasses
    super_arglist = cls.get_super_arglist()
    if super_arglist:
        # Extract superclass names
        superclasses = []
        for child in super_arglist.children:
            if hasattr(child, 'value') and child.value not in (',', ):
                superclasses.append(child.value)
        if superclasses:
            print(f"  Inherits from: {', '.join(superclasses)}")
    else:
        print("  No explicit superclasses")
    
    # Get methods in class
    methods = list(cls.iter_funcdefs())
    print(f"  Methods: {len(methods)}")
    for method in methods:
        print(f"    {method.name.value}()")

# Find lambda expressions
def find_lambdas(node):
    """Find all lambda expressions in a tree."""
    lambdas = []
    if hasattr(node, 'type') and node.type == 'lambdef':
        lambdas.append(node)
    if hasattr(node, 'children'):
        for child in node.children:
            lambdas.extend(find_lambdas(child))
    return lambdas

lambdas = find_lambdas(module)
print(f"\nFound {len(lambdas)} lambda expressions")
for lam in lambdas:
    params = lam.get_params()
    param_names = [p.name.value for p in params]
    print(f"  lambda {', '.join(param_names)}: ...")

class Param(PythonBaseNode):
    """
    Function parameter node.
    
    Attributes:
        type (str): 'param'
        name (Name): Parameter name node
    """
    
    @property
    def star_count(self):
        """
        Number of stars in parameter.
        
        Returns:
            int: 0 for normal, 1 for *args, 2 for **kwargs
        """
    
    @property
    def default(self):
        """
        Parameter default value.
        
        Returns:
            NodeOrLeaf | None: Default value expression or None
        """
    
    @property
    def annotation(self):
        """
        Parameter type annotation.
        
        Returns:
            NodeOrLeaf | None: Type annotation or None
        """
    
    @property
    def position_index(self):
        """
        Positional index of parameter.
        
        Returns:
            int: Position index in parameter list
        """
    
    def get_parent_function(self):
        """
        Get the function containing this parameter.
        
        Returns:
            Function | Lambda: Parent function or lambda
        """

import parso

code = '''
def complex_func(pos_only, /, normal, *args, kw_only, **kwargs):
    pass

def annotated_func(a: int, b: str = "default", c: list[int] = None):
    pass
'''

module = parso.parse(code)

for func in module.iter_funcdefs():
    print(f"\nFunction: {func.name.value}")
    params = func.get_params()
    
    for i, param in enumerate(params):
        print(f"  Parameter {i}: {param.name.value}")
        print(f"    Position index: {param.position_index}")
        print(f"    Star count: {param.star_count}")
        
        if param.annotation:
            print(f"    Annotation: {param.annotation.get_code()}")
        
        if param.default:
            print(f"    Default: {param.default.get_code()}")
        
        parent_func = param.get_parent_function()
        print(f"    Parent function: {parent_func.name.value}")

class Name(PythonLeaf):
    """
    Identifier/name node.
    
    Attributes:
        type (str): 'name'
        value (str): The identifier string
    """
    
    def is_definition(self, include_setitem=False):
        """
        Check if this name is being defined.
        
        Args:
            include_setitem (bool): Include item assignments like obj[key] = value
            
        Returns:
            bool: True if this is a definition
        """
    
    def get_definition(self, import_name_always=False, include_setitem=False):
        """
        Get the definition node for this name.
        
        Args:
            import_name_always (bool): Always consider import names as definitions
            include_setitem (bool): Include item assignments
            
        Returns:
            NodeOrLeaf | None: Definition node or None
        """

import parso

code = '''
import os
from sys import path

x = 42
def func(param):
    global x
    x = param
    local_var = "hello"
    return local_var

class MyClass:
    def __init__(self):
        self.attr = "value"
'''

module = parso.parse(code)

# Analyze all name usages
used_names = module.get_used_names()

for name_str, name_nodes in used_names.items():
    print(f"\nName '{name_str}' appears {len(name_nodes)} times:")
    
    for name_node in name_nodes:
        print(f"  At {name_node.start_pos}: ", end="")
        
        if name_node.is_definition():
            print("DEFINITION", end="")
            definition = name_node.get_definition()
            if definition:
                print(f" (in {definition.type})", end="")
        else:
            print("USAGE", end="")
            definition = name_node.get_definition()
            if definition:
                print(f" (defined in {definition.type})", end="")
        
        print()

# Focus on specific names
if 'x' in used_names:
    print(f"\nVariable 'x' analysis:")
    for x_node in used_names['x']:
        if x_node.is_definition():
            def_node = x_node.get_definition()
            print(f"  Defined at {x_node.start_pos} in {def_node.type}")
        else:
            print(f"  Used at {x_node.start_pos}")

class String(Literal):
    """
    String literal node.
    
    Attributes:
        type (str): 'string'
        value (str): Complete string including quotes and prefix
    """
    
    @property
    def string_prefix(self):
        """
        Get string prefix (r, f, b, u, etc.).
        
        Returns:
            str: String prefix characters
        """

class Number(Literal):
    """
    Numeric literal node.
    
    Attributes:
        type (str): 'number'
        value (str): Number as string (preserves format)
    """

class Keyword(PythonLeaf):
    """
    Python keyword node.
    
    Attributes:
        type (str): 'keyword'
        value (str): Keyword string
    """

class Operator(PythonLeaf):
    """
    Operator node.
    
    Attributes:
        type (str): 'operator'
        value (str): Operator string
    """

import parso

code = '''
x = 42
y = 3.14
z = "hello"
w = r"raw string"
f_str = f"formatted {x}"
binary = 0b1010
hex_num = 0xFF

if x > 0:
    print("positive")
'''

module = parso.parse(code)

def analyze_literals(node):
    """Find and analyze all literals in the tree."""
    literals = {'strings': [], 'numbers': [], 'keywords': [], 'operators': []}
    
    def walk(n):
        if hasattr(n, 'type'):
            if n.type == 'string':
                prefix = n.string_prefix if hasattr(n, 'string_prefix') else ''
                literals['strings'].append((n.value, prefix, n.start_pos))
            elif n.type == 'number':
                literals['numbers'].append((n.value, n.start_pos))
            elif n.type == 'keyword':
                literals['keywords'].append((n.value, n.start_pos))
            elif n.type == 'operator':
                literals['operators'].append((n.value, n.start_pos))
        
        if hasattr(n, 'children'):
            for child in n.children:
                walk(child)
    
    walk(node)
    return literals

literals = analyze_literals(module)

print("String literals:")
for value, prefix, pos in literals['strings']:
    print(f"  {value} (prefix: '{prefix}') at {pos}")

print("\nNumber literals:")
for value, pos in literals['numbers']:
    print(f"  {value} at {pos}")

print("\nKeywords:")
for value, pos in literals['keywords']:
    print(f"  {value} at {pos}")

print("\nOperators:")
for value, pos in literals['operators']:
    if value not in (',', ':', '(', ')', '[', ']'):  # Skip common punctuation
        print(f"  {value} at {pos}")

class IfStmt(Flow):
    """
    If statement node.
    
    Attributes:
        type (str): 'if_stmt'
    """
    
    def get_test_nodes(self):
        """
        Get test expressions for if/elif clauses.
        
        Yields:
            NodeOrLeaf: Test expressions
        """
    
    def is_node_after_else(self, node):
        """
        Check if a node is in the else clause.
        
        Args:
            node (NodeOrLeaf): Node to check
            
        Returns:
            bool: True if node is after else
        """

class ForStmt(Flow):
    """
    For loop node.
    
    Attributes:
        type (str): 'for_stmt'
    """
    
    def get_testlist(self):
        """
        Get the iteration target expression.
        
        Returns:
            NodeOrLeaf: Expression being iterated over
        """
    
    def get_defined_names(self, include_setitem=False):
        """
        Get names defined by loop variable.
        
        Args:
            include_setitem (bool): Include item assignments
            
        Returns:
            list[Name]: Names defined in loop target
        """

class WithStmt(Flow):
    """
    With statement node.
    
    Attributes:
        type (str): 'with_stmt'
    """
    
    def get_defined_names(self, include_setitem=False):
        """
        Get names defined in 'as' clauses.
        
        Args:
            include_setitem (bool): Include item assignments
            
        Returns:
            list[Name]: Names defined by 'as' clauses
        """

class TryStmt(Flow):
    """
    Try statement node.
    
    Attributes:
        type (str): 'try_stmt'
    """
    
    def get_except_clause_tests(self):
        """
        Get exception types from except clauses.
        
        Yields:
            NodeOrLeaf | None: Exception type expressions or None for bare except
        """

import parso

code = '''
if x > 0:
    print("positive")
elif x < 0:
    print("negative")
else:
    print("zero")

for item in items:
    process(item)

for i, value in enumerate(data):
    print(f"{i}: {value}")

with open("file.txt") as f:
    content = f.read()

with lock, open("file") as f:
    data = f.read()

try:
    result = risky_operation()
except ValueError as e:
    print(f"ValueError: {e}")
except (TypeError, AttributeError):
    print("Type or attribute error")
except:
    print("Unknown error")
'''

module = parso.parse(code)

# Analyze control flow statements
def find_control_flow(node):
    """Find all control flow statements."""
    control_flow = []
    
    def walk(n):
        if hasattr(n, 'type'):
            if n.type in ('if_stmt', 'for_stmt', 'while_stmt', 'try_stmt', 'with_stmt'):
                control_flow.append(n)
        
        if hasattr(n, 'children'):
            for child in n.children:
                walk(child)
    
    walk(node)
    return control_flow

statements = find_control_flow(module)

for stmt in statements:
    print(f"\n{stmt.type} at line {stmt.start_pos[0]}:")
    
    if stmt.type == 'if_stmt':
        tests = list(stmt.get_test_nodes())
        print(f"  Test conditions: {len(tests)}")
        for i, test in enumerate(tests):
            print(f"    {i+1}: {test.get_code().strip()}")
    
    elif stmt.type == 'for_stmt':
        target = stmt.get_testlist()
        defined_names = stmt.get_defined_names()
        print(f"  Iterating over: {target.get_code().strip()}")
        print(f"  Loop variables: {[n.value for n in defined_names]}")
    
    elif stmt.type == 'with_stmt':
        defined_names = stmt.get_defined_names()
        if defined_names:
            print(f"  Context variables: {[n.value for n in defined_names]}")
    
    elif stmt.type == 'try_stmt':
        except_tests = list(stmt.get_except_clause_tests())
        print(f"  Exception handlers: {len(except_tests)}")
        for i, test in enumerate(except_tests):
            if test is None:
                print(f"    {i+1}: bare except")
            else:
                print(f"    {i+1}: {test.get_code().strip()}")

class ImportName(Import):
    """
    Regular import statement (import foo).
    
    Attributes:
        type (str): 'import_name'
    """
    
    def get_defined_names(self, include_setitem=False):
        """Get names defined by this import."""
    
    def get_paths(self):
        """Get import paths as lists of names."""
    
    def is_nested(self):
        """Check if import defines nested names."""

class ImportFrom(Import):
    """
    From import statement (from foo import bar).
    
    Attributes:
        type (str): 'import_from'
    """
    
    def get_defined_names(self, include_setitem=False):
        """Get names defined by this import."""
    
    def get_from_names(self):
        """Get the module names being imported from."""
    
    def get_paths(self):
        """Get complete import paths."""
    
    @property
    def level(self):
        """
        Get relative import level.
        
        Returns:
            int: Number of dots in relative import
        """
    
    def is_star_import(self):
        """Check if this is a star import."""

import parso

code = '''
import os
import sys
from pathlib import Path
from . import local_module
from ..parent import parent_func
from foo import bar as baz
from typing import List, Dict, Optional
import numpy as np
from collections import *
'''

module = parso.parse(code)

imports = list(module.iter_imports())

for imp in imports:
    print(f"\n{imp.type} at line {imp.start_pos[0]}:")
    
    if imp.type == 'import_name':
        defined = [n.value for n in imp.get_defined_names()]
        paths = imp.get_paths()
        
        print(f"  Defines: {defined}")
        print(f"  Paths: {[[n.value for n in path] for path in paths]}")
        print(f"  Nested: {imp.is_nested()}")
    
    elif imp.type == 'import_from':
        from_names = [n.value for n in imp.get_from_names()]
        defined = [n.value for n in imp.get_defined_names()]
        level = imp.level
        
        print(f"  From: {'.'.join(from_names) if from_names else '(current)'}")
        print(f"  Level: {level} ({'absolute' if level == 0 else 'relative'})")
        print(f"  Defines: {defined}")
        print(f"  Star import: {imp.is_star_import()}")
        
        paths = imp.get_paths()
        print(f"  Full paths: {[[n.value for n in path] for path in paths]}")