add assignment dict handling to class_handler.py

pythonbpf/vmlinux_parser/class_handler.py

@@ -2,9 +2,9 @@ import logging
 from functools import lru_cache
 import importlib
 
-from .assignment_info import AssignmentInfo
+from .assignment_info import AssignmentInfo, AssignmentType
 from .dependency_handler import DependencyHandler
-from .dependency_node import DependencyNode
+from .dependency_node import DependencyNode, Field
 import ctypes
 from typing import Optional, Any, Dict
 
@@ -103,12 +103,21 @@ def process_vmlinux_post_ast(
             else:
                 raise TypeError("Could not get required class and definition")
 
+            # Create a members dictionary for AssignmentInfo
+            members_dict: Dict[str, tuple[str, Field]] = {}
+
             logger.debug(f"Extracted fields for {current_symbol_name}: {field_table}")
             for elem in field_table.items():
                 elem_name, elem_temp_list = elem
                 [elem_type, elem_bitfield_size] = elem_temp_list
                 local_module_name = getattr(elem_type, "__module__", None)
                 new_dep_node.add_field(elem_name, elem_type, ready=False)
+
+                # Store field reference for struct assignment info
+                field_ref = new_dep_node.get_field(elem_name)
+                if field_ref:
+                    members_dict[elem_name] = (elem_name, field_ref)
+
                 if local_module_name == ctypes.__name__:
                     # TODO: need to process pointer to ctype and also CFUNCTYPES here recursively. Current processing is a single dereference
                     new_dep_node.set_field_bitfield_size(elem_name, elem_bitfield_size)
@@ -220,6 +229,7 @@ def process_vmlinux_post_ast(
                                     containing_type,
                                     llvm_handler,
                                     handler,
+                                    assignments,  # Pass assignments to recursive call
                                     processing_stack,
                                 )
                                 new_dep_node.set_field_ready(elem_name, True)
@@ -237,7 +247,11 @@ def process_vmlinux_post_ast(
                             else str(elem_type)
                         )
                         process_vmlinux_post_ast(
-                            elem_type, llvm_handler, handler, processing_stack
+                            elem_type,
+                            llvm_handler,
+                            handler,
+                            assignments,
+                            processing_stack,
                         )
                         new_dep_node.set_field_ready(elem_name, True)
                 else:
@@ -245,10 +259,21 @@ def process_vmlinux_post_ast(
                         f"{elem_name} with type {elem_type} from module {module_name} not supported in recursive resolver"
                     )
 
+            # Add struct to assignments dictionary
+            assignments[current_symbol_name] = AssignmentInfo(
+                value_type=AssignmentType.STRUCT,
+                python_type=elem_type_class,
+                value=None,
+                pointer_level=None,
+                signature=None,
+                members=members_dict,
+            )
+            logger.info(f"Added struct assignment info for {current_symbol_name}")
+
     else:
         raise ImportError("UNSUPPORTED Module")
 
-    logging.info(
+    logger.info(
         f"{current_symbol_name} processed and handler readiness {handler.is_ready}"
     )
     return True