add dynamic input hook registration

Signed-off-by: Maria Teresa Ortega <teresa.ortega0903@gmail.com>

Author: teresa-ortega

src/lambkin/core/decorators/benchmark.py

@@ -30,6 +30,7 @@
 
 from lambkin.core.ctx.context import Context
 from lambkin.core.ctx.source import Source
+from lambkin.core.decorators.input import InputRegistry
 
 
 def _parse_options(fn, cli_args):
@@ -70,6 +71,8 @@ def benchmark(variants, num_iterations):
         )
 
     def decorator(fn):
+        inputs = InputRegistry()
+
         @functools.wraps(fn)
         def wrapper(args=None, output_dir=None):
             cli_args = sys.argv[1:] if args is None else args
@@ -85,8 +88,10 @@ def wrapper(args=None, output_dir=None):
                         variant_index=variant_index,
                         output_dir=output_dir,
                     )
+                    inputs.resolve(ctx)
                     fn(ctx)
 
+        wrapper.input = inputs.register
         return wrapper
 
     return decorator

src/lambkin/core/decorators/input.py

@@ -14,29 +14,25 @@
 
 """Input decorator for lambkin."""
 
-import functools
 
+class InputRegistry:
+    """Manages the registration and resolution of input hooks for a benchmark."""
 
-def input(hook):
-    """Wrap a benchmark function to inject an input into ``ctx.inputs``.
+    def __init__(self):
+        """Initialize the InputRegistry."""
+        self._hooks = []
 
-    Uses :func:`inspect.getfile` to read the hook name and injects
-    the return value into ``ctx.inputs.<name>`` before the benchmark runs.
+    def register(self, hook_fn):
+        """Decorator used to register a function as an input provider."""
+        self._hooks.append(hook_fn)
+        return hook_fn
 
-    Parameters
-    ----------
-    hook : callable
-        Input resolver function. Must accept a single ``ctx`` argument
-        and return the input value to inject into ``ctx.inputs``.
-    """
+    def resolve(self, ctx):
+        """Resolve all registered input hooks.
 
-    def decorator(benchmark_fn):
-        @functools.wraps(benchmark_fn)
-        def wrapper(ctx):
-            setattr(ctx.inputs, hook.__name__, hook(ctx))
-            return benchmark_fn(ctx)
-
-        wrapper.__wrapped__ = benchmark_fn
-        return wrapper
-
-    return decorator
+        Executes all registered hooks and dynamically maps their return values
+        to the "ctx.inputs" namespace using the original function's name.
+        """
+        for hook in self._hooks:
+            result = hook(ctx)
+            setattr(ctx.inputs, hook.__name__, result)