add assignment info class family and change how assignments are handled

Author: varun-r-mallya

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.

pythonbpf/vmlinux_parser/import_detector.py

@@ -4,6 +4,7 @@
 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")

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()