Apply comments from review.

scripts/core_graph.py

@@ -17,7 +17,8 @@
 from transactron.graph import TracingFragment  # noqa: E402
 from test.test_core import CoreTestElaboratable  # noqa: E402
 from coreblocks.params.configurations import basic_core_config  # noqa: E402
-from transactron.core import TransactionManagerKey, TransactionModule  # noqa: E402
+from transactron.core import TransactionModule  # noqa: E402
+from transactron.core.keys import TransactionManagerKey  # noqa: E402
 
 gp = GenParams(basic_core_config)
 elaboratable = CoreTestElaboratable(gp)

test/regression/pysim.py

@@ -5,7 +5,7 @@
 from amaranth.utils import exact_log2
 from amaranth import *
 
-from transactron.core import TransactionManagerKey
+from transactron.core.keys import TransactionManagerKey
 
 from .memory import *
 from .common import SimulationBackend, SimulationExecutionResult

test/transactions/test_transactions.py

@@ -15,11 +15,8 @@
 from transactron.lib import Adapter, AdapterTrans
 from transactron.utils import Scheduler
 
-from transactron.core import (
-    Priority,
-    trivial_roundrobin_cc_scheduler,
-    eager_deterministic_cc_scheduler,
-)
+from transactron.core import Priority
+from transactron.core.schedulers import 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,5 +1,4 @@
 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

transactron/core/keys.py

@@ -0,0 +1,13 @@
+from transactron.utils import *
+from typing import TYPE_CHECKING
+from dataclasses import dataclass
+
+if TYPE_CHECKING:
+    from .manager import TransactionManager  # noqa: F401 because of https://github.com/PyCQA/pyflakes/issues/571
+
+__all__ = ["TransactionManagerKey"]
+
+
+@dataclass(frozen=True)
+class TransactionManagerKey(SimpleKey["TransactionManager"]):
+    pass

transactron/core/transaction.py

@@ -3,19 +3,14 @@
 from amaranth import tracer
 from typing import Optional, Iterator, TYPE_CHECKING
 from .transaction_base import *
+from .keys import *
 from contextlib import contextmanager
-from dataclasses import dataclass
 
 if TYPE_CHECKING:
     from .tmodule import TModule
     from .manager import TransactionManager
 
-__all__ = ["Transaction", "TransactionManagerKey"]
-
-
-@dataclass(frozen=True)
-class TransactionManagerKey(SimpleKey["TransactionManager"]):
-    pass
+__all__ = ["Transaction"]
 
 
 class Transaction(TransactionBase):

transactron/testing/infrastructure.py

@@ -17,7 +17,8 @@
 from .gtkw_extension import write_vcd_ext
 from transactron import Method
 from transactron.lib import AdapterTrans
-from transactron.core import TransactionManagerKey, TransactionModule
+from transactron.core.keys import TransactionManagerKey
+from transactron.core import TransactionModule
 from transactron.utils import ModuleConnector, HasElaborate, auto_debug_signals, HasDebugSignals
 
 T = TypeVar("T")

transactron/utils/gen.py

@@ -5,7 +5,8 @@
 from amaranth.back import verilog
 from amaranth.hdl import Fragment
 
-from transactron.core import TransactionManager, TransactionManagerKey
+from transactron.core import TransactionManager
+from transactron.core.keys import TransactionManagerKey
 from transactron.core.manager import MethodMap
 from transactron.lib.metrics import HardwareMetricsManager
 from transactron.utils.dependencies import DependencyContext