From e636fcaea788ecec6b1e69f54ee2c45eb48584d8 Mon Sep 17 00:00:00 2001
From: varun-r-mallya <varunrmallya@gmail.com>
Date: Mon, 20 Oct 2025 04:21:42 +0530
Subject: [PATCH] add assignment info class family and change how assignments
 are handled

---
 pythonbpf/vmlinux_parser/assignment_info.py   | 34 +++++++++++++++++++
 pythonbpf/vmlinux_parser/import_detector.py   | 26 +++++++++++---
 .../vmlinux/simple_struct_test.py             |  1 -
 3 files changed, 55 insertions(+), 6 deletions(-)
 create mode 100644 pythonbpf/vmlinux_parser/assignment_info.py

diff --git a/pythonbpf/vmlinux_parser/assignment_info.py b/pythonbpf/vmlinux_parser/assignment_info.py
new file mode 100644
index 0000000..d3665c0
--- /dev/null
+++ b/pythonbpf/vmlinux_parser/assignment_info.py
@@ -0,0 +1,34 @@
+from enum import Enum, auto
+from typing import Any, Callable, Dict, List, Optional, TypedDict
+from dataclasses import dataclass
+
+from pythonbpf.vmlinux_parser.dependency_node import Field
+
+
+@dataclass
+class AssignmentType(Enum):
+    CONSTANT = auto()
+    STRUCT = auto()
+    ARRAY = auto()  # probably won't be used
+    FUNCTION_POINTER = auto()
+    POINTER = auto()  # again, probably won't be used
+
+@dataclass
+class FunctionSignature(TypedDict):
+    return_type: str
+    param_types: List[str]
+    varargs: bool
+
+
+# Thew name of the assignment will be in the dict that uses this class
+@dataclass
+class AssignmentInfo(TypedDict):
+    value_type: AssignmentType
+    python_type: type
+    value: Optional[Any]
+    pointer_level: Optional[int]
+    signature: Optional[FunctionSignature]  # For function pointers
+    # The key of the dict is the name of the field.
+    #   Value is a tuple that contains the global variable representing that field
+    #   along with all the information about that field as a Field type.
+    members: Optional[Dict[str, tuple[str, Field]]]  # For structs.
diff --git a/pythonbpf/vmlinux_parser/import_detector.py b/pythonbpf/vmlinux_parser/import_detector.py
index 972b1ff..965bca8 100644
--- a/pythonbpf/vmlinux_parser/import_detector.py
+++ b/pythonbpf/vmlinux_parser/import_detector.py
@@ -4,6 +4,7 @@ from typing import List, Tuple, Any
 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
@@ -82,7 +83,7 @@ def vmlinux_proc(tree: ast.AST, module):
     # initialise dependency handler
     handler = DependencyHandler()
     # initialise assignment dictionary of name to type
-    assignments: dict[str, tuple[type, Any]] = {}
+    assignments: dict[str, AssignmentInfo] = {}
 
     if not import_statements:
         logger.info("No vmlinux imports found")
@@ -132,16 +133,31 @@ def vmlinux_proc(tree: ast.AST, module):
     return assignments
 
 
-def process_vmlinux_assign(node, module, assignments: dict[str, tuple[type, Any]]):
-    # Check if this is a simple assignment with a constant value
+def process_vmlinux_assign(node, module, assignments: dict[str, AssignmentInfo]):
+    """Process assignments from vmlinux module."""
+    # Only handle single-target assignments
     if len(node.targets) == 1 and isinstance(node.targets[0], ast.Name):
         target_name = node.targets[0].id
+
+        # Handle constant value assignments
         if isinstance(node.value, ast.Constant):
-            assignments[target_name] = (type(node.value.value), node.value.value)
+            # 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
+            )
             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:
-            raise ValueError(f"Unsupported assignment type for {target_name}")
+            logger.warning(
+                f"Unsupported assignment type for {target_name}: {ast.dump(node.value)}"
+            )
     else:
         raise ValueError("Not a simple assignment")
diff --git a/tests/passing_tests/vmlinux/simple_struct_test.py b/tests/passing_tests/vmlinux/simple_struct_test.py
index f47076f..43e5c9e 100644
--- a/tests/passing_tests/vmlinux/simple_struct_test.py
+++ b/tests/passing_tests/vmlinux/simple_struct_test.py
@@ -27,4 +27,3 @@ def LICENSE() -> str:
 
 
 compile_to_ir("simple_struct_test.py", "simple_struct_test.ll")
-compile()