From 8da50b7068f3979e7b8fc3c593f9f6c80866e83e Mon Sep 17 00:00:00 2001
From: varun-r-mallya <varunrmallya@gmail.com>
Date: Mon, 20 Oct 2025 04:31:23 +0530
Subject: [PATCH] float assignments to class_handler.py

---
 pythonbpf/vmlinux_parser/assignment_info.py   |  3 ++-
 pythonbpf/vmlinux_parser/class_handler.py     | 19 +++++++++++++++----
 pythonbpf/vmlinux_parser/import_detector.py   |  7 +++----
 .../vmlinux/simple_struct_test.py             |  2 +-
 4 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/pythonbpf/vmlinux_parser/assignment_info.py b/pythonbpf/vmlinux_parser/assignment_info.py
index d3665c0..435a7ad 100644
--- a/pythonbpf/vmlinux_parser/assignment_info.py
+++ b/pythonbpf/vmlinux_parser/assignment_info.py
@@ -1,5 +1,5 @@
 from enum import Enum, auto
-from typing import Any, Callable, Dict, List, Optional, TypedDict
+from typing import Any, Dict, List, Optional, TypedDict
 from dataclasses import dataclass
 
 from pythonbpf.vmlinux_parser.dependency_node import Field
@@ -13,6 +13,7 @@ class AssignmentType(Enum):
     FUNCTION_POINTER = auto()
     POINTER = auto()  # again, probably won't be used
 
+
 @dataclass
 class FunctionSignature(TypedDict):
     return_type: str
diff --git a/pythonbpf/vmlinux_parser/class_handler.py b/pythonbpf/vmlinux_parser/class_handler.py
index 108fa9f..2adf76d 100644
--- a/pythonbpf/vmlinux_parser/class_handler.py
+++ b/pythonbpf/vmlinux_parser/class_handler.py
@@ -1,6 +1,8 @@
 import logging
 from functools import lru_cache
 import importlib
+
+from .assignment_info import AssignmentInfo
 from .dependency_handler import DependencyHandler
 from .dependency_node import DependencyNode
 import ctypes
@@ -15,17 +17,26 @@ def get_module_symbols(module_name: str):
     return [name for name in dir(imported_module)], imported_module
 
 
-def process_vmlinux_class(node, llvm_module, handler: DependencyHandler):
+def process_vmlinux_class(
+    node,
+    llvm_module,
+    handler: DependencyHandler,
+    assignments: dict[str, AssignmentInfo],
+):
     symbols_in_module, imported_module = get_module_symbols("vmlinux")
     if node.name in symbols_in_module:
         vmlinux_type = getattr(imported_module, node.name)
-        process_vmlinux_post_ast(vmlinux_type, llvm_module, handler)
+        process_vmlinux_post_ast(vmlinux_type, llvm_module, handler, assignments)
     else:
         raise ImportError(f"{node.name} not in vmlinux")
 
 
 def process_vmlinux_post_ast(
-    elem_type_class, llvm_handler, handler: DependencyHandler, processing_stack=None
+    elem_type_class,
+    llvm_handler,
+    handler: DependencyHandler,
+    assignments: dict[str, AssignmentInfo],
+    processing_stack=None,
 ):
     # Initialize processing stack on first call
     if processing_stack is None:
@@ -46,7 +57,7 @@ def process_vmlinux_post_ast(
         logger.debug(f"Node {current_symbol_name} already processed and ready")
         return True
 
-    # XXX:Check it's use. It's probably not being used.
+    # XXX:Check its use. It's probably not being used.
     if current_symbol_name in processing_stack:
         logger.debug(
             f"Dependency already in processing stack for {current_symbol_name}, skipping"
diff --git a/pythonbpf/vmlinux_parser/import_detector.py b/pythonbpf/vmlinux_parser/import_detector.py
index 965bca8..d8bd78f 100644
--- a/pythonbpf/vmlinux_parser/import_detector.py
+++ b/pythonbpf/vmlinux_parser/import_detector.py
@@ -1,6 +1,5 @@
 import ast
 import logging
-from typing import List, Tuple, Any
 import importlib
 import inspect
 
@@ -12,7 +11,7 @@ from .class_handler import process_vmlinux_class
 logger = logging.getLogger(__name__)
 
 
-def detect_import_statement(tree: ast.AST) -> List[Tuple[str, ast.ImportFrom]]:
+def detect_import_statement(tree: ast.AST) -> list[tuple[str, ast.ImportFrom]]:
     """
     Parse AST and detect import statements from vmlinux.
 
@@ -113,7 +112,7 @@ def vmlinux_proc(tree: ast.AST, module):
                     isinstance(mod_node, ast.ClassDef)
                     and mod_node.name == imported_name
                 ):
-                    process_vmlinux_class(mod_node, module, handler)
+                    process_vmlinux_class(mod_node, module, handler, assignments)
                     found = True
                     break
                 if isinstance(mod_node, ast.Assign):
@@ -148,7 +147,7 @@ def process_vmlinux_assign(node, module, assignments: dict[str, AssignmentInfo])
                 value=node.value.value,
                 pointer_level=None,
                 signature=None,
-                members=None
+                members=None,
             )
             logger.info(
                 f"Added assignment: {target_name} = {node.value.value!r} of type {type(node.value.value)}"
diff --git a/tests/passing_tests/vmlinux/simple_struct_test.py b/tests/passing_tests/vmlinux/simple_struct_test.py
index 43e5c9e..c9390c8 100644
--- a/tests/passing_tests/vmlinux/simple_struct_test.py
+++ b/tests/passing_tests/vmlinux/simple_struct_test.py
@@ -1,4 +1,4 @@
-from pythonbpf import bpf, section, bpfglobal, compile_to_ir, compile
+from pythonbpf import bpf, section, bpfglobal, compile_to_ir
 from vmlinux import TASK_COMM_LEN  # noqa: F401
 from vmlinux import struct_trace_event_raw_sys_enter  # noqa: F401