From 97ad2716bd6ef448a661cfa57d01b39dec204656 Mon Sep 17 00:00:00 2001 From: whitequark 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