add multi imports to single import line.

This commit is contained in:
2025-11-08 18:08:04 +05:30
parent a9d82d40d3
commit c5bef26b88
2 changed files with 50 additions and 41 deletions

View File

@ -25,7 +25,7 @@ def detect_import_statement(tree: ast.AST) -> list[tuple[str, ast.ImportFrom]]:
List of tuples containing (module_name, imported_item) for each vmlinux import List of tuples containing (module_name, imported_item) for each vmlinux import
Raises: Raises:
SyntaxError: If multiple imports from vmlinux are attempted or import * is used SyntaxError: If import * is used
""" """
vmlinux_imports = [] vmlinux_imports = []
@ -40,28 +40,19 @@ def detect_import_statement(tree: ast.AST) -> list[tuple[str, ast.ImportFrom]]:
"Please import specific types explicitly." "Please import specific types explicitly."
) )
# Check for multiple imports: from vmlinux import A, B, C
if len(node.names) > 1:
imported_names = [alias.name for alias in node.names]
raise SyntaxError(
f"Multiple imports from vmlinux are not supported. "
f"Found: {', '.join(imported_names)}. "
f"Please use separate import statements for each type."
)
# Check if no specific import is specified (should not happen with valid Python) # Check if no specific import is specified (should not happen with valid Python)
if len(node.names) == 0: if len(node.names) == 0:
raise SyntaxError( raise SyntaxError(
"Import from vmlinux must specify at least one type." "Import from vmlinux must specify at least one type."
) )
# Valid single import # Support multiple imports: from vmlinux import A, B, C
for alias in node.names: for alias in node.names:
import_name = alias.name import_name = alias.name
# Use alias if provided, otherwise use the original name (commented) # Use alias if provided, otherwise use the original name
# as_name = alias.asname if alias.asname else alias.name as_name = alias.asname if alias.asname else alias.name
vmlinux_imports.append(("vmlinux", node)) vmlinux_imports.append(("vmlinux", node, import_name, as_name))
logger.info(f"Found vmlinux import: {import_name}") logger.info(f"Found vmlinux import: {import_name} as {as_name}")
# Handle "import vmlinux" statements (not typical but should be rejected) # Handle "import vmlinux" statements (not typical but should be rejected)
elif isinstance(node, ast.Import): elif isinstance(node, ast.Import):
@ -103,40 +94,37 @@ def vmlinux_proc(tree: ast.AST, module):
with open(source_file, "r") as f: with open(source_file, "r") as f:
mod_ast = ast.parse(f.read(), filename=source_file) mod_ast = ast.parse(f.read(), filename=source_file)
for import_mod, import_node in import_statements: for import_mod, import_node, imported_name, as_name in import_statements:
for alias in import_node.names: found = False
imported_name = alias.name for mod_node in mod_ast.body:
found = False if isinstance(mod_node, ast.ClassDef) and mod_node.name == imported_name:
for mod_node in mod_ast.body: process_vmlinux_class(mod_node, module, handler)
if ( found = True
isinstance(mod_node, ast.ClassDef) break
and mod_node.name == imported_name if isinstance(mod_node, ast.Assign):
): for target in mod_node.targets:
process_vmlinux_class(mod_node, module, handler) if isinstance(target, ast.Name) and target.id == imported_name:
found = True process_vmlinux_assign(mod_node, module, assignments, as_name)
break found = True
if isinstance(mod_node, ast.Assign): break
for target in mod_node.targets: if found:
if isinstance(target, ast.Name) and target.id == imported_name: break
process_vmlinux_assign(mod_node, module, assignments) if not found:
found = True logger.info(f"{imported_name} not found as ClassDef or Assign in vmlinux")
break
if found:
break
if not found:
logger.info(
f"{imported_name} not found as ClassDef or Assign in vmlinux"
)
IRGenerator(module, handler, assignments) IRGenerator(module, handler, assignments)
return assignments return assignments
def process_vmlinux_assign(node, module, assignments: dict[str, AssignmentInfo]): def process_vmlinux_assign(
node, module, assignments: dict[str, AssignmentInfo], target_name=None
):
"""Process assignments from vmlinux module.""" """Process assignments from vmlinux module."""
# Only handle single-target assignments # Only handle single-target assignments
if len(node.targets) == 1 and isinstance(node.targets[0], ast.Name): if len(node.targets) == 1 and isinstance(node.targets[0], ast.Name):
target_name = node.targets[0].id # 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 # Handle constant value assignments
if isinstance(node.value, ast.Constant): if isinstance(node.value, ast.Constant):

View File

@ -0,0 +1,21 @@
from vmlinux import struct_request, struct_pt_regs, XDP_PASS
from pythonbpf import bpf, section, bpfglobal, compile_to_ir
import logging
@bpf
@section("kprobe/blk_mq_start_request")
def example(ctx: struct_pt_regs):
req = struct_request(ctx.di)
c = req.__data_len
d = XDP_PASS
print(f"data length {c} and test {d}")
@bpf
@bpfglobal
def LICENSE() -> str:
return "GPL"
compile_to_ir("requests.py", "requests.ll", loglevel=logging.INFO)