Allow additional method arguments in Adapter (#558)

test/stages/test_retirement.py

@@ -33,7 +33,9 @@ def elaborate(self, platform):
             scheduler_layouts.free_rf_layout, 2**self.gen_params.phys_regs_bits
         )
 
-        m.submodules.mock_rob_peek = self.mock_rob_peek = TestbenchIO(Adapter(o=rob_layouts.peek_layout))
+        m.submodules.mock_rob_peek = self.mock_rob_peek = TestbenchIO(
+            Adapter(o=rob_layouts.peek_layout, nonexclusive=True)
+        )
 
         m.submodules.mock_rob_retire = self.mock_rob_retire = TestbenchIO(Adapter(o=rob_layouts.retire_layout))
 

transactron/lib/adapters.py

@@ -3,7 +3,7 @@
 from ..utils import SrcLoc, get_src_loc
 from ..core import *
 from ..core import SignalBundle
-from typing import Optional
+from ..utils._typing import type_self_kwargs_as
 
 __all__ = [
     "AdapterBase",
@@ -92,22 +92,19 @@ class Adapter(AdapterBase):
         Data passed as argument to the defined method.
     """
 
-    def __init__(
-        self, *, name: Optional[str] = None, i: MethodLayout = (), o: MethodLayout = (), src_loc: int | SrcLoc = 0
-    ):
+    @type_self_kwargs_as(Method.__init__)
+    def __init__(self, **kwargs):
         """
         Parameters
         ----------
-        i: record layout
-            The input layout of the defined method.
-        o: record layout
-            The output layout of the defined method.
-        src_loc: int | SrcLoc
-            How many stack frames deep the source location is taken from.
-            Alternatively, the source location to use instead of the default.
+        **kwargs
+            Keyword arguments for Method that will be created.
+            See transactron.core.Method.__init__ for parameters description.
         """
-        src_loc = get_src_loc(src_loc)
-        super().__init__(Method(name=name, i=i, o=o, src_loc=src_loc))
+
+        kwargs["src_loc"] = get_src_loc(kwargs.setdefault("src_loc", 0))
+
+        super().__init__(Method(**kwargs))
         self.data_in = Record.like(self.iface.data_out)
         self.data_out = Record.like(self.iface.data_in)
 

transactron/utils/_typing.py

@@ -1,10 +1,14 @@
 from typing import (
+    Callable,
+    Concatenate,
     Generic,
     NoReturn,
     Optional,
+    ParamSpec,
     Protocol,
     TypeAlias,
     TypeVar,
+    cast,
     runtime_checkable,
     Union,
     Any,
@@ -64,6 +68,7 @@
 
 T = TypeVar("T")
 U = TypeVar("U")
+P = ParamSpec("P")
 
 ROGraph: TypeAlias = Mapping[T, Iterable[T]]
 Graph: TypeAlias = dict[T, set[T]]
@@ -136,3 +141,16 @@ def elaborate(self, platform) -> "HasElaborate":
 class HasDebugSignals(Protocol):
     def debug_signals(self) -> SignalBundle:
         ...
+
+
+def type_self_kwargs_as(as_func: Callable[Concatenate[Any, P], Any]):
+    """
+    Decorator used to annotate `**kwargs` type to be the same as named arguments from `as_func` method.
+
+    Works only with methods with (self, **kwargs) signature. `self` parameter is also required in `as_func`.
+    """
+
+    def return_func(func: Callable[Concatenate[Any, ...], T]) -> Callable[Concatenate[Any, P], T]:
+        return cast(Callable[Concatenate[Any, P], T], func)
+
+    return return_func