From 97ad2716bd6ef448a661cfa57d01b39dec204656 Mon Sep 17 00:00:00 2001
From: whitequark <cz@m-labs.hk>
Date: Sun, 22 Sep 2019 15:16:36 +0000
Subject: [PATCH] hdl.rec: allow using Enum subclass as shape.

Fixes #223.
---
 nmigen/hdl/rec.py           |  3 +++
 nmigen/test/test_hdl_rec.py | 14 ++++++++++++++
 2 files changed, 17 insertions(+)

diff --git a/nmigen/hdl/rec.py b/nmigen/hdl/rec.py
index d3d1b33..53c08e7 100644
--- 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 2491587..03d29e6 100644
--- 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),
-- 
2.30.2