Consistently use Dataclass syntac for AssignmentInfo and related classes

Author: r41k0u

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,15 +15,15 @@ class AssignmentType(Enum):
 
 
 @dataclass
-class FunctionSignature(TypedDict):
+class FunctionSignature:
     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):
+class AssignmentInfo:
     value_type: AssignmentType
     python_type: type
     value: Optional[Any]

pythonbpf/vmlinux_parser/vmlinux_exports_handler.py

@@ -37,38 +37,38 @@ def is_vmlinux_enum(self, name):
         """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")
 
     def is_vmlinux_struct(self, name):
         """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
 
     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 @@ def handle_vmlinux_struct_field(
     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 @@ def get_field_type(self, vmlinux_struct_name, field_name):
     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(