From 7bd882958b0ce4627a7c01546fe6ec30c2277539 Mon Sep 17 00:00:00 2001
From: Lekcyjna <309016@uwr.edu.pl>
Date: Tue, 12 Mar 2024 18:23:41 +0100
Subject: [PATCH] Fix imports.

---
 coreblocks/structs_common/rs.py           |  2 +-
 test/transactions/test_branches.py        |  4 +-
 test/transactions/test_transaction_lib.py |  3 +-
 test/transactions/test_transactions.py    |  2 +-
 transactron/core.py                       | 25 ----------
 transactron/core/__init__.py              | 14 +++---
 transactron/core/manager.py               | 59 ++++++++++++++++++++--
 transactron/core/method.py                |  7 +--
 transactron/core/schedulers.py            |  5 +-
 transactron/core/sugar.py                 | 60 ++---------------------
 transactron/core/tmodule.py               |  3 ++
 transactron/core/transaction.py           | 17 ++++---
 transactron/core/transaction_base.py      | 10 ++--
 transactron/lib/adapters.py               |  3 +-
 transactron/lib/buttons.py                |  2 +-
 transactron/lib/connectors.py             |  2 +-
 transactron/lib/reqres.py                 |  2 +-
 transactron/lib/storage.py                |  2 +-
 transactron/lib/transformers.py           | 12 ++++-
 transactron/profiler.py                   |  3 +-
 transactron/testing/profiler.py           |  3 +-
 transactron/testing/testbenchio.py        |  3 +-
 transactron/utils/_typing.py              |  2 +
 transactron/utils/gen.py                  |  3 +-
 24 files changed, 118 insertions(+), 130 deletions(-)
 delete mode 100644 transactron/core.py

diff --git a/coreblocks/structs_common/rs.py b/coreblocks/structs_common/rs.py
index fe8d04ba4..6af0b5e2a 100644
--- a/coreblocks/structs_common/rs.py
+++ b/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"]
diff --git a/test/transactions/test_branches.py b/test/transactions/test_branches.py
index f66b954b7..2b38636a3 100644
--- a/test/transactions/test_branches.py
+++ b/test/transactions/test_branches.py
@@ -1,8 +1,6 @@
 from amaranth import *
 from itertools import product
 from transactron.core import (
-    CtrlPath,
-    MethodMap,
     TModule,
     Method,
     Transaction,
@@ -10,6 +8,8 @@
     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
diff --git a/test/transactions/test_transaction_lib.py b/test/transactions/test_transaction_lib.py
index 058557f22..dd3899964 100644
--- a/test/transactions/test_transaction_lib.py
+++ b/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,
diff --git a/test/transactions/test_transactions.py b/test/transactions/test_transactions.py
index c73a1642b..1f8240aab 100644
--- a/test/transactions/test_transactions.py
+++ b/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
 
 
diff --git a/transactron/core.py b/transactron/core.py
deleted file mode 100644
index 69a456927..000000000
--- a/transactron/core.py
+++ /dev/null
@@ -1,25 +0,0 @@
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
-
diff --git a/transactron/core/__init__.py b/transactron/core/__init__.py
index ebe4a9a53..5f46c9c4e 100644
--- a/transactron/core/__init__.py
+++ b/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
diff --git a/transactron/core/manager.py b/transactron/core/manager.py
index 428473fa0..0b3c55090 100644
--- a/transactron/core/manager.py
+++ b/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
diff --git a/transactron/core/method.py b/transactron/core/method.py
index 2180c29c0..98fb59f3d 100644
--- a/transactron/core/method.py
+++ b/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.
 
diff --git a/transactron/core/schedulers.py b/transactron/core/schedulers.py
index b4cc32464..856d4450b 100644
--- a/transactron/core/schedulers.py
+++ b/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
 
diff --git a/transactron/core/sugar.py b/transactron/core/sugar.py
index 960f71498..49ca6e5cd 100644
--- a/transactron/core/sugar.py
+++ b/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,
 ):
diff --git a/transactron/core/tmodule.py b/transactron/core/tmodule.py
index 192d31cfa..a9ee030b1 100644
--- a/transactron/core/tmodule.py
+++ b/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.
 
diff --git a/transactron/core/transaction.py b/transactron/core/transaction.py
index 00f4d1f0e..5568a448d 100644
--- a/transactron/core/transaction.py
+++ b/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
diff --git a/transactron/core/transaction_base.py b/transactron/core/transaction_base.py
index cb49757a0..be4fe6f93 100644
--- a/transactron/core/transaction_base.py
+++ b/transactron/core/transaction_base.py
@@ -4,7 +4,6 @@
 from itertools import count
 from typing import (
     ClassVar,
-    NoReturn,
     TypeAlias,
     TypedDict,
     Union,
@@ -16,7 +15,6 @@
     Iterator,
 )
 from amaranth import *
-from dataclasses import dataclass
 
 from .tmodule import TModule, CtrlPath
 from transactron.graph import Owned
@@ -25,16 +23,12 @@
 if TYPE_CHECKING:
     from .method import Method
     from .transaction import Transaction
-    from .manager import TransactionManager
 
-__all__ = ["TransactionBase", "TransactionManagerKey"]
+__all__ = ["TransactionBase", "Priority"]
 
 TransactionOrMethod: TypeAlias = Union["Transaction", "Method"]
 TransactionOrMethodBound = TypeVar("TransactionOrMethodBound", "Transaction", "Method")
 
-@dataclass(frozen=True)
-class TransactionManagerKey(SimpleKey[TransactionManager]):
-    pass
 
 class Priority(Enum):
     #: Conflicting transactions/methods don't have a priority order.
@@ -44,6 +38,7 @@ class Priority(Enum):
     #: Right transaction/method is prioritized over the left one.
     RIGHT = auto()
 
+
 class RelationBase(TypedDict):
     end: TransactionOrMethod
     priority: Priority
@@ -54,6 +49,7 @@ class RelationBase(TypedDict):
 class Relation(RelationBase):
     start: TransactionOrMethod
 
+
 @runtime_checkable
 class TransactionBase(Owned, Protocol):
     stack: ClassVar[list[Union["Transaction", "Method"]]] = []
diff --git a/transactron/lib/adapters.py b/transactron/lib/adapters.py
index ed7f2640f..4a7ea8460 100644
--- a/transactron/lib/adapters.py
+++ b/transactron/lib/adapters.py
@@ -2,8 +2,7 @@
 
 from ..utils import SrcLoc, get_src_loc, MethodStruct
 from ..core import *
-from ..core import SignalBundle
-from ..utils._typing import type_self_kwargs_as
+from ..utils._typing import type_self_kwargs_as, SignalBundle
 
 __all__ = [
     "AdapterBase",
diff --git a/transactron/lib/buttons.py b/transactron/lib/buttons.py
index 59bf081b5..d275cd25d 100644
--- a/transactron/lib/buttons.py
+++ b/transactron/lib/buttons.py
@@ -2,7 +2,7 @@
 
 from transactron.utils.transactron_helpers import from_method_layout
 from ..core import *
-from ..utils import SrcLoc, get_src_loc
+from ..utils import SrcLoc, get_src_loc, MethodLayout
 
 __all__ = ["ClickIn", "ClickOut"]
 
diff --git a/transactron/lib/connectors.py b/transactron/lib/connectors.py
index b9a6eb204..511cf6248 100644
--- a/transactron/lib/connectors.py
+++ b/transactron/lib/connectors.py
@@ -3,7 +3,7 @@
 
 from transactron.utils.transactron_helpers import from_method_layout
 from ..core import *
-from ..utils import SrcLoc, get_src_loc
+from ..utils import SrcLoc, get_src_loc, MethodLayout
 
 __all__ = [
     "FIFO",
diff --git a/transactron/lib/reqres.py b/transactron/lib/reqres.py
index 518d53443..f9aeb6e06 100644
--- a/transactron/lib/reqres.py
+++ b/transactron/lib/reqres.py
@@ -1,6 +1,6 @@
 from amaranth import *
 from ..core import *
-from ..utils import SrcLoc, get_src_loc
+from ..utils import SrcLoc, get_src_loc, MethodLayout
 from .connectors import Forwarder, FIFO
 from transactron.lib import BasicFifo
 from amaranth.utils import *
diff --git a/transactron/lib/storage.py b/transactron/lib/storage.py
index 9549527f8..22e46f2d1 100644
--- a/transactron/lib/storage.py
+++ b/transactron/lib/storage.py
@@ -3,7 +3,7 @@
 
 from transactron.utils.transactron_helpers import from_method_layout
 from ..core import *
-from ..utils import SrcLoc, get_src_loc
+from ..utils import SrcLoc, get_src_loc, MethodLayout
 from typing import Optional
 from transactron.utils import assign, AssignType, LayoutList
 from .reqres import ArgumentsToResultsZipper
diff --git a/transactron/lib/transformers.py b/transactron/lib/transformers.py
index a1445fcf5..f874cea2c 100644
--- a/transactron/lib/transformers.py
+++ b/transactron/lib/transformers.py
@@ -2,11 +2,19 @@
 
 from transactron.utils.transactron_helpers import get_src_loc
 from ..core import *
-from ..core import RecordDict
 from ..utils import SrcLoc
 from typing import Optional, Protocol
 from collections.abc import Callable
-from transactron.utils import ValueLike, assign, AssignType, ModuleLike, MethodStruct, HasElaborate
+from transactron.utils import (
+    ValueLike,
+    assign,
+    AssignType,
+    ModuleLike,
+    MethodStruct,
+    HasElaborate,
+    MethodLayout,
+    RecordDict,
+)
 from .connectors import Forwarder, ManyToOneConnectTrans, ConnectTrans
 from .simultaneous import condition
 
diff --git a/transactron/profiler.py b/transactron/profiler.py
index 0132b2ef7..28574731a 100644
--- a/transactron/profiler.py
+++ b/transactron/profiler.py
@@ -4,7 +4,8 @@
 from dataclasses import dataclass, field
 from dataclasses_json import dataclass_json
 from transactron.utils import SrcLoc, IdGenerator
-from transactron.core import MethodMap, TransactionManager
+from transactron.core import TransactionManager
+from transactron.core.manager import MethodMap
 
 
 __all__ = [
diff --git a/transactron/testing/profiler.py b/transactron/testing/profiler.py
index 18451112c..b70119c7d 100644
--- a/transactron/testing/profiler.py
+++ b/transactron/testing/profiler.py
@@ -1,5 +1,6 @@
 from amaranth.sim import *
-from transactron.core import MethodMap, TransactionManager
+from transactron.core import TransactionManager
+from transactron.core.manager import MethodMap
 from transactron.profiler import CycleProfile, MethodSamples, Profile, ProfileData, ProfileSamples, TransactionSamples
 from .functions import TestGen
 
diff --git a/transactron/testing/testbenchio.py b/transactron/testing/testbenchio.py
index 8f9cc253d..cbe4c446f 100644
--- a/transactron/testing/testbenchio.py
+++ b/transactron/testing/testbenchio.py
@@ -2,8 +2,7 @@
 from amaranth.sim import Settle, Passive
 from typing import Optional, Callable
 from transactron.lib import AdapterBase
-from transactron.core import ValueLike, SignalBundle
-from transactron.utils import mock_def_helper
+from transactron.utils import ValueLike, SignalBundle, mock_def_helper
 from transactron.utils._typing import RecordIntDictRet, RecordValueDict, RecordIntDict
 from .functions import set_inputs, get_outputs, TestGen
 
diff --git a/transactron/utils/_typing.py b/transactron/utils/_typing.py
index 8f42c1910..e529eee24 100644
--- a/transactron/utils/_typing.py
+++ b/transactron/utils/_typing.py
@@ -45,6 +45,7 @@
     "RecordIntDict",
     "RecordIntDictRet",
     "RecordValueDict",
+    "RecordDict",
     "ROGraph",
     "Graph",
     "GraphCC",
@@ -76,6 +77,7 @@
 RecordIntDict: TypeAlias = Mapping[str, Union[int, "RecordIntDict"]]
 RecordIntDictRet: TypeAlias = Mapping[str, Any]  # full typing hard to work with
 RecordValueDict: TypeAlias = Mapping[str, Union[ValueLike, "RecordValueDict"]]
+RecordDict: TypeAlias = ValueLike | Mapping[str, "RecordDict"]
 
 T = TypeVar("T")
 U = TypeVar("U")
diff --git a/transactron/utils/gen.py b/transactron/utils/gen.py
index 2ff40dec2..e76bc359e 100644
--- a/transactron/utils/gen.py
+++ b/transactron/utils/gen.py
@@ -5,7 +5,8 @@
 from amaranth.back import verilog
 from amaranth.hdl import Fragment
 
-from transactron.core import TransactionManager, MethodMap, TransactionManagerKey
+from transactron.core import TransactionManager, TransactionManagerKey
+from transactron.core.manager import MethodMap
 from transactron.lib.metrics import HardwareMetricsManager
 from transactron.utils.dependencies import DependencyContext
 from transactron.utils.idgen import IdGenerator