Fix imports.

coreblocks/structs_common/rs.py

@@ -4,7 +4,7 @@
 from amaranth.lib.coding import PriorityEncoder
 from transactron import Method, def_method, TModule
 from coreblocks.params import RSLayouts, GenParams, OpType
-from transactron.core import RecordDict
+from transactron.utils import RecordDict
 from transactron.utils.transactron_helpers import make_layout
 
 __all__ = ["RS"]

test/transactions/test_branches.py

@@ -1,15 +1,15 @@
 from amaranth import *
 from itertools import product
 from transactron.core import (
-    CtrlPath,
-    MethodMap,
     TModule,
     Method,
     Transaction,
     TransactionManager,
     TransactionModule,
     def_method,
 )
+from transactron.core.tmodule import CtrlPath
+from transactron.core.manager import MethodMap
 from unittest import TestCase
 from transactron.testing import TestCaseWithSimulator
 from transactron.utils.dependencies import DependencyContext

test/transactions/test_transaction_lib.py

@@ -10,10 +10,9 @@
 
 from amaranth import *
 from transactron import *
-from transactron.core import RecordDict
 from transactron.lib import *
 from coreblocks.utils import *
-from transactron.utils._typing import ModuleLike, MethodStruct
+from transactron.utils._typing import ModuleLike, MethodStruct, RecordDict
 from transactron.utils import ModuleConnector
 from transactron.testing import (
     SimpleTestCircuit,

test/transactions/test_transactions.py

@@ -17,10 +17,10 @@
 
 from transactron.core import (
     Priority,
-    TransactionScheduler,
     trivial_roundrobin_cc_scheduler,
     eager_deterministic_cc_scheduler,
 )
+from transactron.core.manager import TransactionScheduler
 from transactron.utils.dependencies import DependencyContext
 
 

transactron/core/__init__.py

@@ -1,7 +1,7 @@
-from .tmodule import *
-from .schedulers import *
-from .transaction_base import *
-from .method import *
-from .transaction import*
-from .manager import *
-from .sugar import *
+from .tmodule import *  # noqa: F401
+from .schedulers import *  # noqa: F401
+from .transaction_base import *  # noqa: F401
+from .method import *  # noqa: F401
+from .transaction import *  # noqa: F401
+from .manager import *  # noqa: F401
+from .sugar import *  # noqa: F401

transactron/core/manager.py

@@ -2,7 +2,6 @@
 from typing import Callable, Iterable, Sequence, TypeAlias, Tuple
 from os import environ
 from graphlib import TopologicalSorter
-from dataclasses import dataclass
 from amaranth import *
 from itertools import chain, filterfalse, product
 
@@ -12,17 +11,16 @@
 
 from .transaction_base import TransactionBase, TransactionOrMethod, Priority, Relation
 from .method import Method
-from .transaction import Transaction
+from .transaction import Transaction, TransactionManagerKey
 from .tmodule import TModule
 from .schedulers import eager_deterministic_cc_scheduler
 
-__all__ =  ["MethodMap", "TransactionManager"] 
+__all__ = ["TransactionManager", "TransactionModule"]
 
 TransactionGraph: TypeAlias = Graph["Transaction"]
 TransactionGraphCC: TypeAlias = GraphCC["Transaction"]
 PriorityOrder: TypeAlias = dict["Transaction", int]
 TransactionScheduler: TypeAlias = Callable[["MethodMap", TransactionGraph, TransactionGraphCC, PriorityOrder], Module]
-RecordDict: TypeAlias = ValueLike | Mapping[str, "RecordDict"]
 
 
 class MethodMap:
@@ -69,6 +67,7 @@ def transactions(self) -> Collection["Transaction"]:
     def methods_and_transactions(self) -> Iterable[TransactionOrMethod]:
         return chain(self.methods, self.transactions)
 
+
 class TransactionManager(Elaboratable):
     """Transaction manager
 
@@ -433,3 +432,55 @@ def method_debug(m: Method):
             "transactions": {t.name: transaction_debug(t) for t in method_map.transactions},
             "methods": {m.owned_name: method_debug(m) for m in method_map.methods},
         }
+
+
+class TransactionModule(Elaboratable):
+    """
+    `TransactionModule` is used as wrapper on `Elaboratable` classes,
+    which adds support for transactions. It creates a
+    `TransactionManager` which will handle transaction scheduling
+    and can be used in definition of `Method`\\s and `Transaction`\\s.
+    The `TransactionManager` is stored in a `DependencyManager`.
+    """
+
+    def __init__(
+        self,
+        elaboratable: HasElaborate,
+        dependency_manager: Optional[DependencyManager] = None,
+        transaction_manager: Optional[TransactionManager] = None,
+    ):
+        """
+        Parameters
+        ----------
+        elaboratable: HasElaborate
+            The `Elaboratable` which should be wrapped to add support for
+            transactions and methods.
+        dependency_manager: DependencyManager, optional
+            The `DependencyManager` to use inside the transaction module.
+            If omitted, a new one is created.
+        transaction_manager: TransactionManager, optional
+            The `TransactionManager` to use inside the transaction module.
+            If omitted, a new one is created.
+        """
+        if transaction_manager is None:
+            transaction_manager = TransactionManager()
+        if dependency_manager is None:
+            dependency_manager = DependencyManager()
+        self.manager = dependency_manager
+        self.manager.add_dependency(TransactionManagerKey(), transaction_manager)
+        self.elaboratable = elaboratable
+
+    def context(self) -> DependencyContext:
+        return DependencyContext(self.manager)
+
+    def elaborate(self, platform):
+        with silence_mustuse(self.manager.get_dependency(TransactionManagerKey())):
+            with self.context():
+                elaboratable = Fragment.get(self.elaboratable, platform)
+
+        m = Module()
+
+        m.submodules.main_module = elaboratable
+        m.submodules.transactionManager = self.manager.get_dependency(TransactionManagerKey())
+
+        return m

transactron/core/method.py

@@ -12,6 +12,7 @@
 
 __all__ = ["Method"]
 
+
 class Method(TransactionBase):
     """Transactional method.
 
@@ -129,7 +130,7 @@ def like(other: "Method", *, name: Optional[str] = None, src_loc: int | SrcLoc =
         """
         return Method(name=name, i=other.layout_in, o=other.layout_out, src_loc=get_src_loc(src_loc))
 
-    def proxy(self, m: 'TModule', method: "Method"):
+    def proxy(self, m: "TModule", method: "Method"):
         """Define as a proxy for another method.
 
         The calls to this method will be forwarded to `method`.
@@ -150,7 +151,7 @@ def _(arg):
     @contextmanager
     def body(
         self,
-        m: 'TModule',
+        m: "TModule",
         *,
         ready: ValueLike = C(1),
         out: ValueLike = C(0, 0),
@@ -219,7 +220,7 @@ def _validate_arguments(self, arg_rec: MethodStruct) -> ValueLike:
         return self.ready
 
     def __call__(
-        self, m: 'TModule', arg: Optional[AssignArg] = None, enable: ValueLike = C(1), /, **kwargs: AssignArg
+        self, m: "TModule", arg: Optional[AssignArg] = None, enable: ValueLike = C(1), /, **kwargs: AssignArg
     ) -> MethodStruct:
         """Call a method.
 

transactron/core/schedulers.py

@@ -7,8 +7,9 @@
 
 __all__ = ["eager_deterministic_cc_scheduler", "trivial_roundrobin_cc_scheduler"]
 
+
 def eager_deterministic_cc_scheduler(
-    method_map: 'MethodMap', gr: 'TransactionGraph', cc: 'TransactionGraphCC', porder: 'PriorityOrder'
+    method_map: "MethodMap", gr: "TransactionGraph", cc: "TransactionGraphCC", porder: "PriorityOrder"
 ) -> Module:
     """eager_deterministic_cc_scheduler
 
@@ -44,7 +45,7 @@ def eager_deterministic_cc_scheduler(
 
 
 def trivial_roundrobin_cc_scheduler(
-    method_map: 'MethodMap', gr: 'TransactionGraph', cc: 'TransactionGraphCC', porder: 'PriorityOrder'
+    method_map: "MethodMap", gr: "TransactionGraph", cc: "TransactionGraphCC", porder: "PriorityOrder"
 ) -> Module:
     """trivial_roundrobin_cc_scheduler
 

transactron/core/sugar.py

@@ -3,70 +3,16 @@
 from transactron.utils import *
 from transactron.utils.assign import AssignArg
 
-from .manager import TransactionManager
-from .transaction_base import TransactionManagerKey
-
 if TYPE_CHECKING:
     from .tmodule import TModule
     from .method import Method
 
-__all__ = ["TransactionModule","def_method"]
-
-class TransactionModule(Elaboratable):
-    """
-    `TransactionModule` is used as wrapper on `Elaboratable` classes,
-    which adds support for transactions. It creates a
-    `TransactionManager` which will handle transaction scheduling
-    and can be used in definition of `Method`\\s and `Transaction`\\s.
-    The `TransactionManager` is stored in a `DependencyManager`.
-    """
-
-    def __init__(
-        self,
-        elaboratable: HasElaborate,
-        dependency_manager: Optional[DependencyManager] = None,
-        transaction_manager: Optional[TransactionManager] = None,
-    ):
-        """
-        Parameters
-        ----------
-        elaboratable: HasElaborate
-            The `Elaboratable` which should be wrapped to add support for
-            transactions and methods.
-        dependency_manager: DependencyManager, optional
-            The `DependencyManager` to use inside the transaction module.
-            If omitted, a new one is created.
-        transaction_manager: TransactionManager, optional
-            The `TransactionManager` to use inside the transaction module.
-            If omitted, a new one is created.
-        """
-        if transaction_manager is None:
-            transaction_manager = TransactionManager()
-        if dependency_manager is None:
-            dependency_manager = DependencyManager()
-        self.manager = dependency_manager
-        self.manager.add_dependency(TransactionManagerKey(), transaction_manager)
-        self.elaboratable = elaboratable
-
-    def context(self) -> DependencyContext:
-        return DependencyContext(self.manager)
-
-    def elaborate(self, platform):
-        with silence_mustuse(self.manager.get_dependency(TransactionManagerKey())):
-            with self.context():
-                elaboratable = Fragment.get(self.elaboratable, platform)
-
-        m = Module()
-
-        m.submodules.main_module = elaboratable
-        m.submodules.transactionManager = self.manager.get_dependency(TransactionManagerKey())
-
-        return m
+__all__ = ["def_method"]
 
 
 def def_method(
-    m: 'TModule',
-    method: 'Method',
+    m: "TModule",
+    method: "Method",
     ready: ValueLike = C(1),
     validate_arguments: Optional[Callable[..., ValueLike]] = None,
 ):

transactron/core/tmodule.py

@@ -8,6 +8,7 @@
 
 __all__ = ["TModule"]
 
+
 class _AvoidingModuleBuilderDomain:
     """
     A wrapper over Amaranth domain to abstract away internal Amaranth implementation.
@@ -46,6 +47,7 @@ def __setattr__(self, name: str, value):
     def __setitem__(self, name: str, value):
         return self.__setattr__(name, value)
 
+
 class EnterType(Enum):
     """Characterizes stack behavior of Amaranth's context managers for control structures."""
 
@@ -157,6 +159,7 @@ def build_ctrl_path(self):
         """Returns the current control path."""
         return CtrlPath(self.module, self.ctrl_path[:])
 
+
 class TModule(ModuleLike, Elaboratable):
     """Extended Amaranth module for use with transactions.
 

transactron/core/transaction.py

@@ -1,17 +1,22 @@
 from transactron.utils import *
 from amaranth import *
 from amaranth import tracer
-from typing import Optional, Callable, Iterator, TYPE_CHECKING
+from typing import Optional, Iterator, TYPE_CHECKING
 from .transaction_base import *
-from .sugar import def_method
 from contextlib import contextmanager
-from transactron.utils.assign import AssignArg
+from dataclasses import dataclass
 
 if TYPE_CHECKING:
     from .tmodule import TModule
     from .manager import TransactionManager
 
-__all__ = ["Transaction"]
+__all__ = ["Transaction", "TransactionManagerKey"]
+
+
+@dataclass(frozen=True)
+class TransactionManagerKey(SimpleKey["TransactionManager"]):
+    pass
+
 
 class Transaction(TransactionBase):
     """Transaction.
@@ -52,7 +57,7 @@ class Transaction(TransactionBase):
     """
 
     def __init__(
-        self, *, name: Optional[str] = None, manager: Optional['TransactionManager'] = None, src_loc: int | SrcLoc = 0
+        self, *, name: Optional[str] = None, manager: Optional["TransactionManager"] = None, src_loc: int | SrcLoc = 0
     ):
         """
         Parameters
@@ -80,7 +85,7 @@ def __init__(
         self.grant = Signal(name=self.owned_name + "_grant")
 
     @contextmanager
-    def body(self, m: TModule, *, request: ValueLike = C(1)) -> Iterator["Transaction"]:
+    def body(self, m: "TModule", *, request: ValueLike = C(1)) -> Iterator["Transaction"]:
         """Defines the `Transaction` body.
 
         This context manager allows to conveniently define the actions