From 33e18f6d6dcb1b2dd84229766b56214559c37019 Mon Sep 17 00:00:00 2001
From: Pragyansh Chaturvedi <pragyanshchaturvedi18@gmail.com>
Date: Mon, 3 Nov 2025 21:21:13 +0530
Subject: [PATCH] Introduce HelperSignature in HelperHandlerRegistry

---
 pythonbpf/helper/helper_registry.py | 23 +++++++++++++++++++----
 1 file changed, 19 insertions(+), 4 deletions(-)

diff --git a/pythonbpf/helper/helper_registry.py b/pythonbpf/helper/helper_registry.py
index 476e3b6..7fbb597 100644
--- a/pythonbpf/helper/helper_registry.py
+++ b/pythonbpf/helper/helper_registry.py
@@ -1,17 +1,31 @@
+from dataclasses import dataclass
+from llvmlite import ir
 from typing import Callable
 
 
+@dataclass
+class HelperSignature:
+    """Signature of a BPF helper function"""
+
+    arg_types: list[ir.Type]
+    return_type: ir.Type
+    func: Callable
+
+
 class HelperHandlerRegistry:
     """Registry for BPF helpers"""
 
-    _handlers: dict[str, Callable] = {}
+    _handlers: dict[str, HelperSignature] = {}
 
     @classmethod
-    def register(cls, helper_name):
+    def register(cls, helper_name, param_types=None, return_type=None):
         """Decorator to register a handler function for a helper"""
 
         def decorator(func):
-            cls._handlers[helper_name] = func
+            helper_sig = HelperSignature(
+                arg_types=param_types, return_type=return_type, func=func
+            )
+            cls._handlers[helper_name] = helper_sig
             return func
 
         return decorator
@@ -19,7 +33,8 @@ class HelperHandlerRegistry:
     @classmethod
     def get_handler(cls, helper_name):
         """Get the handler function for a helper"""
-        return cls._handlers.get(helper_name)
+        handler = cls._handlers.get(helper_name)
+        return handler.func if handler else None
 
     @classmethod
     def has_handler(cls, helper_name):