From 96216d441158e5e1fc5fc93468d5d52d1d008f18 Mon Sep 17 00:00:00 2001
From: Pragyansh Chaturvedi <pragyanshchaturvedi18@gmail.com>
Date: Sat, 25 Oct 2025 05:10:47 +0530
Subject: [PATCH] Consistently use Dataclass syntac for AssignmentInfo and
 related classes

---
 pythonbpf/vmlinux_parser/assignment_info.py   |  7 +++---
 .../vmlinux_parser/vmlinux_exports_handler.py | 22 +++++++++----------
 2 files changed, 14 insertions(+), 15 deletions(-)

diff --git a/pythonbpf/vmlinux_parser/assignment_info.py b/pythonbpf/vmlinux_parser/assignment_info.py
index 465432d..1092978 100644
--- a/pythonbpf/vmlinux_parser/assignment_info.py
+++ b/pythonbpf/vmlinux_parser/assignment_info.py
@@ -1,12 +1,11 @@
 from enum import Enum, auto
-from typing import Any, Dict, List, Optional, TypedDict
+from typing import Any, Dict, List, Optional
 from dataclasses import dataclass
 import llvmlite.ir as ir
 
 from pythonbpf.vmlinux_parser.dependency_node import Field
 
 
-@dataclass
 class AssignmentType(Enum):
     CONSTANT = auto()
     STRUCT = auto()
@@ -16,7 +15,7 @@ class AssignmentType(Enum):
 
 
 @dataclass
-class FunctionSignature(TypedDict):
+class FunctionSignature:
     return_type: str
     param_types: List[str]
     varargs: bool
@@ -24,7 +23,7 @@ class FunctionSignature(TypedDict):
 
 # Thew name of the assignment will be in the dict that uses this class
 @dataclass
-class AssignmentInfo(TypedDict):
+class AssignmentInfo:
     value_type: AssignmentType
     python_type: type
     value: Optional[Any]
diff --git a/pythonbpf/vmlinux_parser/vmlinux_exports_handler.py b/pythonbpf/vmlinux_parser/vmlinux_exports_handler.py
index d36c31e..3a3164c 100644
--- a/pythonbpf/vmlinux_parser/vmlinux_exports_handler.py
+++ b/pythonbpf/vmlinux_parser/vmlinux_exports_handler.py
@@ -37,16 +37,16 @@ class VmlinuxHandler:
         """Check if name is a vmlinux enum constant"""
         return (
             name in self.vmlinux_symtab
-            and self.vmlinux_symtab[name]["value_type"] == AssignmentType.CONSTANT
+            and self.vmlinux_symtab[name].value_type == AssignmentType.CONSTANT
         )
 
     def get_vmlinux_struct_type(self, name):
         """Check if name is a vmlinux struct type"""
         if (
             name in self.vmlinux_symtab
-            and self.vmlinux_symtab[name]["value_type"] == AssignmentType.STRUCT
+            and self.vmlinux_symtab[name].value_type == AssignmentType.STRUCT
         ):
-            return self.vmlinux_symtab[name]["python_type"]
+            return self.vmlinux_symtab[name].python_type
         else:
             raise ValueError(f"{name} is not a vmlinux struct type")
 
@@ -54,13 +54,13 @@ class VmlinuxHandler:
         """Check if name is a vmlinux struct"""
         return (
             name in self.vmlinux_symtab
-            and self.vmlinux_symtab[name]["value_type"] == AssignmentType.STRUCT
+            and self.vmlinux_symtab[name].value_type == AssignmentType.STRUCT
         )
 
     def handle_vmlinux_enum(self, name):
         """Handle vmlinux enum constants by returning LLVM IR constants"""
         if self.is_vmlinux_enum(name):
-            value = self.vmlinux_symtab[name]["value"]
+            value = self.vmlinux_symtab[name].value
             logger.info(f"Resolving vmlinux enum {name} = {value}")
             return ir.Constant(ir.IntType(64), value), ir.IntType(64)
         return None
@@ -68,7 +68,7 @@ class VmlinuxHandler:
     def get_vmlinux_enum_value(self, name):
         """Handle vmlinux enum constants by returning LLVM IR constants"""
         if self.is_vmlinux_enum(name):
-            value = self.vmlinux_symtab[name]["value"]
+            value = self.vmlinux_symtab[name].value
             logger.info(f"The value of vmlinux enum {name} = {value}")
             return value
         return None
@@ -115,16 +115,16 @@ class VmlinuxHandler:
     def has_field(self, struct_name, field_name):
         """Check if a vmlinux struct has a specific field"""
         if self.is_vmlinux_struct(struct_name):
-            python_type = self.vmlinux_symtab[struct_name]["python_type"]
+            python_type = self.vmlinux_symtab[struct_name].python_type
             return hasattr(python_type, field_name)
         return False
 
     def get_field_type(self, vmlinux_struct_name, field_name):
         """Get the type of a field in a vmlinux struct"""
         if self.is_vmlinux_struct(vmlinux_struct_name):
-            python_type = self.vmlinux_symtab[vmlinux_struct_name]["python_type"]
+            python_type = self.vmlinux_symtab[vmlinux_struct_name].python_type
             if hasattr(python_type, field_name):
-                return self.vmlinux_symtab[vmlinux_struct_name]["members"][field_name]
+                return self.vmlinux_symtab[vmlinux_struct_name].members[field_name]
             else:
                 raise ValueError(
                     f"Field {field_name} not found in vmlinux struct {vmlinux_struct_name}"
@@ -135,10 +135,10 @@ class VmlinuxHandler:
     def get_field_index(self, vmlinux_struct_name, field_name):
         """Get the type of a field in a vmlinux struct"""
         if self.is_vmlinux_struct(vmlinux_struct_name):
-            python_type = self.vmlinux_symtab[vmlinux_struct_name]["python_type"]
+            python_type = self.vmlinux_symtab[vmlinux_struct_name].python_type
             if hasattr(python_type, field_name):
                 return list(
-                    self.vmlinux_symtab[vmlinux_struct_name]["members"].keys()
+                    self.vmlinux_symtab[vmlinux_struct_name].members.keys()
                 ).index(field_name)
             else:
                 raise ValueError(