From 60c7cf4f996a9bbcd906b1446e1d96c666b7eeca Mon Sep 17 00:00:00 2001
From: Lekcyjna <309016@uwr.edu.pl>
Date: Mon, 23 Oct 2023 20:07:20 +0200
Subject: [PATCH] Move types.

---
 coreblocks/utils/_typing.py |  6 ++++++
 test/common/functions.py    |  2 +-
 test/common/sugar.py        |  2 +-
 test/common/testbenchio.py  |  3 ++-
 transactron/_utils.py       | 10 +++-------
 5 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/coreblocks/utils/_typing.py b/coreblocks/utils/_typing.py
index 124066124..f9a2b8c45 100644
--- a/coreblocks/utils/_typing.py
+++ b/coreblocks/utils/_typing.py
@@ -12,6 +12,8 @@
     Mapping,
     TypeVar,
     runtime_checkable,
+    Union,
+    Any,
 )
 from enum import Enum
 from amaranth import *
@@ -32,6 +34,10 @@
 SignalBundle: TypeAlias = Signal | Record | View | Iterable["SignalBundle"] | Mapping[str, "SignalBundle"]
 LayoutList = list[Tuple[str, ShapeLike | "LayoutList"]]
 
+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"]]
+
 
 class _ModuleBuilderDomainsLike(Protocol):
     def __getattr__(self, name: str) -> _ModuleBuilderDomain:
diff --git a/test/common/functions.py b/test/common/functions.py
index cc85722d4..6c7fa4412 100644
--- a/test/common/functions.py
+++ b/test/common/functions.py
@@ -1,6 +1,6 @@
 from amaranth import *
 from typing import TYPE_CHECKING
-from transactron._utils import RecordValueDict, RecordIntDict
+from coreblocks.utils._typing import RecordValueDict, RecordIntDict
 
 
 if TYPE_CHECKING:
diff --git a/test/common/sugar.py b/test/common/sugar.py
index e134199fe..beb4acf3a 100644
--- a/test/common/sugar.py
+++ b/test/common/sugar.py
@@ -1,7 +1,7 @@
 import functools
 from typing import Callable, Any, Optional
 from .testbenchio import TestbenchIO, TestGen
-from transactron._utils import RecordIntDict
+from coreblocks.utils._typing import RecordIntDict
 
 
 def def_method_mock(
diff --git a/test/common/testbenchio.py b/test/common/testbenchio.py
index 82248c4cb..c1101cb79 100644
--- a/test/common/testbenchio.py
+++ b/test/common/testbenchio.py
@@ -5,7 +5,8 @@
 from typing import Any, Generator, TypeVar, Optional, Callable
 from transactron.lib import AdapterBase
 from transactron.core import ValueLike, SignalBundle
-from transactron._utils import method_def_helper, RecordIntDictRet, RecordValueDict, RecordIntDict
+from transactron._utils import method_def_helper
+from coreblocks.utils._typing import RecordIntDictRet, RecordValueDict, RecordIntDict
 from .functions import set_inputs, get_outputs
 
 T = TypeVar("T")
diff --git a/transactron/_utils.py b/transactron/_utils.py
index 774b38af6..84d97f616 100644
--- a/transactron/_utils.py
+++ b/transactron/_utils.py
@@ -1,9 +1,9 @@
 import itertools
 import sys
 from inspect import Parameter, signature
-from typing import Callable, Iterable, Optional, TypeAlias, TypeVar, Mapping, Union, Any
+from typing import Callable, Iterable, Optional, TypeAlias, TypeVar, Mapping
 from amaranth import *
-from coreblocks.utils._typing import LayoutLike, ValueLike
+from coreblocks.utils._typing import LayoutLike, ShapeLike
 from coreblocks.utils import OneHotSwitchDynamic
 
 __all__ = [
@@ -17,10 +17,6 @@
     "method_def_helper",
 ]
 
-RecordIntDict = Mapping[str, Union[int, "RecordIntDict"]]
-RecordIntDictRet = Mapping[str, Any]  # full typing hard to work with
-RecordValueDict = Mapping[str, Union[ValueLike, "RecordValueDict"]]
-
 
 T = TypeVar("T")
 
@@ -156,7 +152,7 @@ def get_caller_class_name(default: Optional[str] = None) -> tuple[Optional[Elabo
         raise RuntimeError("Not called from a method")
 
 
-def data_layout(val: int) -> LayoutLike:
+def data_layout(val: ShapeLike) -> LayoutLike:
     return [("data", val)]