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 53c08e77e5a315d1c0173dc86cdd9159f8b3b684..9ad5ba79b92d0e281aa2326d0b8bd538592fbbe2 100644 (file)
--- 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 03d29e601052d52b4c5316051522aece88f1d8c8..bae4fb1b18c29c1903e9e592ce927b80f735172e 100644 (file)
--- 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")])