From: whitequark <cz@m-labs.hk>
Date: Sun, 22 Sep 2019 17:23:32 +0000 (+0000)
Subject: hdl.rec: fix using Enum subclass as shape if direction is specified.
X-Git-Url: https://git.libre-soc.org/?a=commitdiff_plain;h=5f216a624870f2e77ac2f4dbd6aa0bb8f7383e89;p=nmigen.git

hdl.rec: fix using Enum subclass as shape if direction is specified.

Also improves error messages.

Fixes #224.
---

diff --git a/nmigen/hdl/rec.py b/nmigen/hdl/rec.py
index 53c08e7..9ad5ba7 100644
--- a/nmigen/hdl/rec.py
+++ b/nmigen/hdl/rec.py
@@ -35,8 +35,6 @@ class Layout:
             if len(field) == 2:
                 name, shape = field
                 direction = DIR_NONE
-                if isinstance(shape, type) and issubclass(shape, Enum):
-                    shape = _enum_shape(shape)
                 if isinstance(shape, list):
                     shape = Layout.wrap(shape)
             else:
@@ -48,9 +46,11 @@ class Layout:
             if not isinstance(name, str):
                 raise TypeError("Field {!r} has invalid name: should be a string"
                                 .format(field))
+            if isinstance(shape, type) and issubclass(shape, Enum):
+                shape = _enum_shape(shape)
             if not isinstance(shape, (int, tuple, Layout)):
-                raise TypeError("Field {!r} has invalid shape: should be an int, tuple, or list "
-                                "of fields of a nested record"
+                raise TypeError("Field {!r} has invalid shape: should be an int, tuple, Enum, or "
+                                "list of fields of a nested record"
                                 .format(field))
             if name in self.fields:
                 raise NameError("Field {!r} has a name that is already present in the layout"
diff --git a/nmigen/test/test_hdl_rec.py b/nmigen/test/test_hdl_rec.py
index 03d29e6..bae4fb1 100644
--- a/nmigen/test/test_hdl_rec.py
+++ b/nmigen/test/test_hdl_rec.py
@@ -36,8 +36,10 @@ class LayoutTestCase(FHDLTestCase):
     def test_enum_field(self):
         layout = Layout.wrap([
             ("enum", UnsignedEnum),
+            ("enum_dir", UnsignedEnum, DIR_FANOUT),
         ])
         self.assertEqual(layout["enum"], ((2, False), DIR_NONE))
+        self.assertEqual(layout["enum_dir"], ((2, False), DIR_FANOUT))
 
     def test_slice_tuple(self):
         layout = Layout.wrap([
@@ -75,7 +77,7 @@ class LayoutTestCase(FHDLTestCase):
 
     def test_wrong_shape(self):
         with self.assertRaises(TypeError,
-                msg="Field ('a', 'x') has invalid shape: should be an int, tuple, or "
+                msg="Field ('a', 'x') has invalid shape: should be an int, tuple, Enum, or "
                     "list of fields of a nested record"):
             Layout.wrap([("a", "x")])