hdl.rec: allow using Enum subclass as shape.

Fixes #223.

diff --git a/nmigen/hdl/rec.py b/nmigen/hdl/rec.py
index d3d1b33ecf9137ae9fbf6897e0ec7f0be5a3a5a4..53c08e77e5a315d1c0173dc86cdd9159f8b3b684 100644 (file)
--- a/nmigen/hdl/rec.py
+++ b/nmigen/hdl/rec.py
@@ -5,6 +5,7 @@ from functools import reduce
 from .. import tracer
 from ..tools import union
 from .ast import *
+from .ast import _enum_shape
 
 
 __all__ = ["Direction", "DIR_NONE", "DIR_FANOUT", "DIR_FANIN", "Layout", "Record"]
@@ -34,6 +35,8 @@ 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:

diff --git a/nmigen/test/test_hdl_rec.py b/nmigen/test/test_hdl_rec.py
index 2491587d2924c5ef9cd6d3168d5f4e8d37ce4dc1..03d29e601052d52b4c5316051522aece88f1d8c8 100644 (file)
--- a/nmigen/test/test_hdl_rec.py
+++ b/nmigen/test/test_hdl_rec.py
@@ -1,8 +1,16 @@
+from enum import Enum
+
 from ..hdl.ast import *
 from ..hdl.rec import *
 from .tools import *
 
 
+class UnsignedEnum(Enum):
+    FOO = 1
+    BAR = 2
+    BAZ = 3
+
+
 class LayoutTestCase(FHDLTestCase):
     def test_fields(self):
         layout = Layout.wrap([
@@ -25,6 +33,12 @@ class LayoutTestCase(FHDLTestCase):
         self.assertEqual(sublayout["a"], ((1, False), DIR_NONE))
         self.assertEqual(sublayout["b"], ((1, False), DIR_NONE))
 
+    def test_enum_field(self):
+        layout = Layout.wrap([
+            ("enum", UnsignedEnum),
+        ])
+        self.assertEqual(layout["enum"], ((2, False), DIR_NONE))
+
     def test_slice_tuple(self):
         layout = Layout.wrap([
             ("a", 1),