import ast import logging import importlib import inspect from .assignment_info import AssignmentInfo, AssignmentType from .dependency_handler import DependencyHandler from .ir_gen import IRGenerator from .class_handler import process_vmlinux_class logger = logging.getLogger(__name__) def detect_import_statement( tree: ast.AST, ) -> list[tuple[str, ast.ImportFrom, str, str]]: """ Parse AST and detect import statements from vmlinux. Returns a list of tuples (module_name, imported_item) for vmlinux imports. Raises SyntaxError for invalid import patterns. Args: tree: The AST to parse Returns: List of tuples containing (module_name, imported_item) for each vmlinux import Raises: SyntaxError: If import * is used """ vmlinux_imports = [] for node in ast.walk(tree): # Handle "from vmlinux import ..." statements if isinstance(node, ast.ImportFrom): if node.module == "vmlinux": # Check for wildcard import: from vmlinux import * if any(alias.name == "*" for alias in node.names): raise SyntaxError( "Wildcard imports from vmlinux are not supported. " "Please import specific types explicitly." ) # Check if no specific import is specified (should not happen with valid Python) if len(node.names) == 0: raise SyntaxError( "Import from vmlinux must specify at least one type." ) # Support multiple imports: from vmlinux import A, B, C for alias in node.names: import_name = alias.name # Use alias if provided, otherwise use the original name as_name = alias.asname if alias.asname else alias.name vmlinux_imports.append(("vmlinux", node, import_name, as_name)) logger.info(f"Found vmlinux import: {import_name} as {as_name}") # Handle "import vmlinux" statements (not typical but should be rejected) elif isinstance(node, ast.Import): for alias in node.names: if alias.name == "vmlinux" or alias.name.startswith("vmlinux."): raise SyntaxError( "Direct import of vmlinux module is not supported. " "Use 'from vmlinux import ' instead." ) logger.info(f"Total vmlinux imports detected: {len(vmlinux_imports)}") # print(f"\n**************\n{vmlinux_imports}\n**************\n") return vmlinux_imports def vmlinux_proc(tree: ast.AST, module): import_statements = detect_import_statement(tree) # initialise dependency handler handler = DependencyHandler() # initialise assignment dictionary of name to type assignments: dict[str, AssignmentInfo] = {} if not import_statements: logger.info("No vmlinux imports found") return None # Import vmlinux module directly try: vmlinux_mod = importlib.import_module("vmlinux") except ImportError: logger.warning("Could not import vmlinux module") return None source_file = inspect.getsourcefile(vmlinux_mod) if source_file is None: logger.warning("Cannot find source for vmlinux module") return None with open(source_file, "r") as f: mod_ast = ast.parse(f.read(), filename=source_file) for import_mod, import_node, imported_name, as_name in import_statements: found = False for mod_node in mod_ast.body: if isinstance(mod_node, ast.ClassDef) and mod_node.name == imported_name: process_vmlinux_class(mod_node, module, handler) found = True break if isinstance(mod_node, ast.Assign): for target in mod_node.targets: if isinstance(target, ast.Name) and target.id == imported_name: process_vmlinux_assign(mod_node, module, assignments, as_name) found = True break if found: break if not found: logger.info(f"{imported_name} not found as ClassDef or Assign in vmlinux") IRGenerator(module, handler, assignments) return assignments def process_vmlinux_assign( node, module, assignments: dict[str, AssignmentInfo], target_name=None ): """Process assignments from vmlinux module.""" # Only handle single-target assignments if len(node.targets) == 1 and isinstance(node.targets[0], ast.Name): # Use provided target_name (for aliased imports) or fall back to original name if target_name is None: target_name = node.targets[0].id # Handle constant value assignments if isinstance(node.value, ast.Constant): # Fixed: using proper TypedDict creation syntax with named arguments assignments[target_name] = AssignmentInfo( value_type=AssignmentType.CONSTANT, python_type=type(node.value.value), value=node.value.value, pointer_level=None, signature=None, members=None, debug_info=None, ) logger.info( f"Added assignment: {target_name} = {node.value.value!r} of type {type(node.value.value)}" ) # Handle other assignment types that we may need to support else: logger.warning( f"Unsupported assignment type for {target_name}: {ast.dump(node.value)}" ) else: raise ValueError("Not a simple assignment")