From 88fc55cb0f4fd8b147dfee3e5803464de358c5dc Mon Sep 17 00:00:00 2001
From: Lekcyjna <309016@uwr.edu.pl>
Date: Sun, 17 Mar 2024 11:20:13 +0100
Subject: [PATCH] Apply comments from review.

---
 scripts/core_graph.py                  |  3 ++-
 test/regression/pysim.py               |  2 +-
 test/transactions/test_transactions.py |  7 ++-----
 transactron/core/__init__.py           |  1 -
 transactron/core/keys.py               | 13 +++++++++++++
 transactron/core/transaction.py        |  9 ++-------
 transactron/testing/infrastructure.py  |  3 ++-
 transactron/utils/gen.py               |  3 ++-
 8 files changed, 24 insertions(+), 17 deletions(-)
 create mode 100644 transactron/core/keys.py

diff --git a/scripts/core_graph.py b/scripts/core_graph.py
index 6818f6dd0..5024a88f8 100755
--- a/scripts/core_graph.py
+++ b/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)
diff --git a/test/regression/pysim.py b/test/regression/pysim.py
index 424d83d8e..661e6d56a 100644
--- a/test/regression/pysim.py
+++ b/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
diff --git a/test/transactions/test_transactions.py b/test/transactions/test_transactions.py
index 1f8240aab..53dec0821 100644
--- a/test/transactions/test_transactions.py
+++ b/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
 
diff --git a/transactron/core/__init__.py b/transactron/core/__init__.py
index 5f46c9c4e..6ead593f8 100644
--- a/transactron/core/__init__.py
+++ b/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
diff --git a/transactron/core/keys.py b/transactron/core/keys.py
new file mode 100644
index 000000000..9444dce34
--- /dev/null
+++ b/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
diff --git a/transactron/core/transaction.py b/transactron/core/transaction.py
index 5568a448d..c6f4176ab 100644
--- a/transactron/core/transaction.py
+++ b/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):
diff --git a/transactron/testing/infrastructure.py b/transactron/testing/infrastructure.py
index dc4a5404e..28faca7d6 100644
--- a/transactron/testing/infrastructure.py
+++ b/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")
diff --git a/transactron/utils/gen.py b/transactron/utils/gen.py
index e76bc359e..7df93cb4f 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, 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