fix Record and Layout typing

diff --git a/nmigen/hdl/rec.pyi b/nmigen/hdl/rec.pyi
index 577a5434308bf9441e25b939abc18881ccfe19e4..ef564a7337cacd7635a705cbec8d8e99c09e5ea4 100644 (file)
--- a/nmigen/hdl/rec.pyi
+++ b/nmigen/hdl/rec.pyi
@@ -1,5 +1,5 @@
 import enum
-from typing import List, Union, Tuple, Any, Dict, Optional
+from typing import List, Union, Tuple, Any, Dict, Optional, Generator
 from .ast import Signal, Value
 
 __all__ = ["Direction", "DIR_NONE", "DIR_FANOUT",
@@ -17,58 +17,77 @@ DIR_FANOUT = Direction.FANOUT
 DIR_FANIN = Direction.FANIN
 
 # recursive types are not yet supported by mypy, manually recurse a few times
-LayoutInputFields0 = List[Union[Tuple[str,
-                                      Union[int,
-                                            Tuple[int, bool],
-                                            Any]],
-                                Tuple[str,
-                                      Union[int,
-                                            Tuple[int, bool]],
-                                      Direction]]]
-
-LayoutInputFields1 = List[Union[Tuple[str,
-                                      Union[int,
-                                            Tuple[int, bool],
-                                            LayoutInputFields0]],
-                                Tuple[str,
-                                      Union[int,
-                                            Tuple[int, bool]],
-                                      Direction]]]
-
-LayoutInputFields2 = List[Union[Tuple[str,
-                                      Union[int,
-                                            Tuple[int, bool],
-                                            LayoutInputFields1]],
-                                Tuple[str,
-                                      Union[int,
-                                            Tuple[int, bool]],
-                                      Direction]]]
-
-LayoutInputFields3 = List[Union[Tuple[str,
-                                      Union[int,
-                                            Tuple[int, bool],
-                                            LayoutInputFields2]],
-                                Tuple[str,
-                                      Union[int,
-                                            Tuple[int, bool]],
-                                      Direction]]]
-
-LayoutInputFields = List[Union[Tuple[str,
-                                     Union[int,
-                                           Tuple[int, bool],
-                                           LayoutInputFields3]],
-                               Tuple[str,
-                                     Union[int,
-                                           Tuple[int, bool]],
-                                     Direction]]]
+LayoutInputFields0 = Union['Layout',
+                           List[Union[Tuple[str,
+                                            Union[int,
+                                                  Tuple[int, bool],
+                                                  Any]],
+                                      Tuple[str,
+                                            Union[int,
+                                                  Tuple[int, bool]],
+                                            Direction]]]]
+
+LayoutInputFields1 = Union['Layout',
+                           List[Union[Tuple[str,
+                                            Union[int,
+                                                  Tuple[int, bool],
+                                                  LayoutInputFields0]],
+                                      Tuple[str,
+                                            Union[int,
+                                                  Tuple[int, bool]],
+                                            Direction]]]]
+
+LayoutInputFields2 = Union['Layout',
+                           List[Union[Tuple[str,
+                                            Union[int,
+                                                  Tuple[int, bool],
+                                                  LayoutInputFields1]],
+                                      Tuple[str,
+                                            Union[int,
+                                                  Tuple[int, bool]],
+                                            Direction]]]]
+
+LayoutInputFields3 = Union['Layout',
+                           List[Union[Tuple[str,
+                                            Union[int,
+                                                  Tuple[int, bool],
+                                                  LayoutInputFields2]],
+                                      Tuple[str,
+                                            Union[int,
+                                                  Tuple[int, bool]],
+                                            Direction]]]]
+
+LayoutInputFields = Union['Layout',
+                          List[Union[Tuple[str,
+                                           Union[int,
+                                                 Tuple[int, bool],
+                                                 LayoutInputFields3]],
+                                     Tuple[str,
+                                           Union[int,
+                                                 Tuple[int, bool]],
+                                           Direction]]]]
+
+LayoutFieldShape = Union[int, Tuple[int, bool], 'Layout']
 
 
 class Layout:
-    fields: Dict[str, Tuple[Union[int, Tuple[int, bool], 'Layout'], Direction]]
+    fields: Dict[str, Tuple[LayoutFieldShape, Direction]]
+
+    @staticmethod
+    def wrap(obj: LayoutInputFields) -> 'Layout':
+        ...
 
     def __init__(self, fields: LayoutInputFields) -> None:
         ...
 
+    def __getitem__(self, name: str) -> Tuple[LayoutFieldShape, Direction]:
+        ...
+
+    def __iter__(self) -> Generator[Tuple[str, LayoutFieldShape, Direction],
+                                    None,
+                                    None]:
+        ...
+
 
 class Record(Value):
     layout: Layout