Consistently use Dataclass syntac for AssignmentInfo and related classes

This commit is contained in:
Pragyansh Chaturvedi
2025-10-25 05:10:47 +05:30
parent 028d9c2c08
commit 96216d4411
2 changed files with 14 additions and 15 deletions

View File

@ -1,12 +1,11 @@
from enum import Enum, auto from enum import Enum, auto
from typing import Any, Dict, List, Optional, TypedDict from typing import Any, Dict, List, Optional
from dataclasses import dataclass from dataclasses import dataclass
import llvmlite.ir as ir import llvmlite.ir as ir
from pythonbpf.vmlinux_parser.dependency_node import Field from pythonbpf.vmlinux_parser.dependency_node import Field
@dataclass
class AssignmentType(Enum): class AssignmentType(Enum):
CONSTANT = auto() CONSTANT = auto()
STRUCT = auto() STRUCT = auto()
@ -16,7 +15,7 @@ class AssignmentType(Enum):
@dataclass @dataclass
class FunctionSignature(TypedDict): class FunctionSignature:
return_type: str return_type: str
param_types: List[str] param_types: List[str]
varargs: bool varargs: bool
@ -24,7 +23,7 @@ class FunctionSignature(TypedDict):
# Thew name of the assignment will be in the dict that uses this class # Thew name of the assignment will be in the dict that uses this class
@dataclass @dataclass
class AssignmentInfo(TypedDict): class AssignmentInfo:
value_type: AssignmentType value_type: AssignmentType
python_type: type python_type: type
value: Optional[Any] value: Optional[Any]

View File

@ -37,16 +37,16 @@ class VmlinuxHandler:
"""Check if name is a vmlinux enum constant""" """Check if name is a vmlinux enum constant"""
return ( return (
name in self.vmlinux_symtab 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): def get_vmlinux_struct_type(self, name):
"""Check if name is a vmlinux struct type""" """Check if name is a vmlinux struct type"""
if ( if (
name in self.vmlinux_symtab 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: else:
raise ValueError(f"{name} is not a vmlinux struct type") raise ValueError(f"{name} is not a vmlinux struct type")
@ -54,13 +54,13 @@ class VmlinuxHandler:
"""Check if name is a vmlinux struct""" """Check if name is a vmlinux struct"""
return ( return (
name in self.vmlinux_symtab 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): def handle_vmlinux_enum(self, name):
"""Handle vmlinux enum constants by returning LLVM IR constants""" """Handle vmlinux enum constants by returning LLVM IR constants"""
if self.is_vmlinux_enum(name): 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}") logger.info(f"Resolving vmlinux enum {name} = {value}")
return ir.Constant(ir.IntType(64), value), ir.IntType(64) return ir.Constant(ir.IntType(64), value), ir.IntType(64)
return None return None
@ -68,7 +68,7 @@ class VmlinuxHandler:
def get_vmlinux_enum_value(self, name): def get_vmlinux_enum_value(self, name):
"""Handle vmlinux enum constants by returning LLVM IR constants""" """Handle vmlinux enum constants by returning LLVM IR constants"""
if self.is_vmlinux_enum(name): 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}") logger.info(f"The value of vmlinux enum {name} = {value}")
return value return value
return None return None
@ -115,16 +115,16 @@ class VmlinuxHandler:
def has_field(self, struct_name, field_name): def has_field(self, struct_name, field_name):
"""Check if a vmlinux struct has a specific field""" """Check if a vmlinux struct has a specific field"""
if self.is_vmlinux_struct(struct_name): 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 hasattr(python_type, field_name)
return False return False
def get_field_type(self, vmlinux_struct_name, field_name): def get_field_type(self, vmlinux_struct_name, field_name):
"""Get the type of a field in a vmlinux struct""" """Get the type of a field in a vmlinux struct"""
if self.is_vmlinux_struct(vmlinux_struct_name): 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): 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: else:
raise ValueError( raise ValueError(
f"Field {field_name} not found in vmlinux struct {vmlinux_struct_name}" 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): def get_field_index(self, vmlinux_struct_name, field_name):
"""Get the type of a field in a vmlinux struct""" """Get the type of a field in a vmlinux struct"""
if self.is_vmlinux_struct(vmlinux_struct_name): 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): if hasattr(python_type, field_name):
return list( return list(
self.vmlinux_symtab[vmlinux_struct_name]["members"].keys() self.vmlinux_symtab[vmlinux_struct_name].members.keys()
).index(field_name) ).index(field_name)
else: else:
raise ValueError( raise ValueError(