Repository: amaranth-lang/amaranth
Pull request #1322 (2024-04-11 17:29:08)

hdl._ast: add enum name argument to `Format.Enum`.
Turns out that RTLIL enum representation requires such, so add a place to store it.


Patch:
diff --git a/amaranth/hdl/_ast.py b/amaranth/hdl/_ast.py
index 4b56d07..86abc42 100644
--- a/amaranth/hdl/_ast.py
+++ b/amaranth/hdl/_ast.py
@@ -2769,8 +2769,11 @@ class Format(_FormatLike):
 
 
     class Enum(_FormatLike):
-        def __init__(self, value, /, variants):
+        def __init__(self, value, /, variants, *, name=None):
             self._value = Value.cast(value)
+            if name is not None and not isinstance(name, str):
+                raise TypeError(f"Enum name must be a string or None, not {name!r}")
+            self._name = name
             if isinstance(variants, EnumMeta):
                 self._variants = {Const.cast(member.value).value: member.name for member in variants}
             else:
@@ -2796,7 +2799,8 @@ class Format(_FormatLike):
                 f" ({val!r} {name!r})"
                 for val, name in self._variants.items()
             )
-            return f"(format-enum {self._value!r}{variants})"
+            name = "-" if self._name is None else repr(self._name)
+            return f"(format-enum {self._value!r} {name}{variants})"
 
 
     class Struct(_FormatLike):
diff --git a/amaranth/lib/enum.py b/amaranth/lib/enum.py
index 8026760..0e2bdea 100644
--- a/amaranth/lib/enum.py
+++ b/amaranth/lib/enum.py
@@ -179,7 +179,7 @@ class EnumType(ShapeCastable, py_enum.EnumMeta):
     def format(cls, value, format_spec):
         if format_spec != "":
             raise ValueError(f"Format specifier {format_spec!r} is not supported for enums")
-        return Format.Enum(value, cls)
+        return Format.Enum(value, cls, name=cls.__name__)
 
     def _value_repr(cls, value):
         yield Repr(FormatEnum(cls), value)


Test patch:
diff --git a/tests/test_hdl_ast.py b/tests/test_hdl_ast.py
index d39da7f..515eb65 100644
--- a/tests/test_hdl_ast.py
+++ b/tests/test_hdl_ast.py
@@ -1740,7 +1740,7 @@ class FormatEnumTestCase(FHDLTestCase):
     def test_construct(self):
         a = Signal(3)
         fmt = Format.Enum(a, {1: "A", 2: "B", 3: "C"})
-        self.assertRepr(fmt, "(format-enum (sig a) (1 'A') (2 'B') (3 'C'))")
+        self.assertRepr(fmt, "(format-enum (sig a) - (1 'A') (2 'B') (3 'C'))")
         self.assertRepr(Format("{}", fmt), """
             (format '{:s}' (switch-value (sig a)
                 (case 001 (const 8'd65))
@@ -1755,8 +1755,8 @@ class FormatEnumTestCase(FHDLTestCase):
             B = 3
             C = 4
 
-        fmt = Format.Enum(a, MyEnum)
-        self.assertRepr(fmt, "(format-enum (sig a) (0 'A') (3 'B') (4 'C'))")
+        fmt = Format.Enum(a, MyEnum, name="MyEnum")
+        self.assertRepr(fmt, "(format-enum (sig a) 'MyEnum' (0 'A') (3 'B') (4 'C'))")
         self.assertRepr(Format("{}", fmt), """
             (format '{:s}' (switch-value (sig a)
                 (case 000 (const 8'd65))
@@ -1774,6 +1774,9 @@ class FormatEnumTestCase(FHDLTestCase):
         with self.assertRaisesRegex(TypeError,
                 r"^Variant names must be strings, not 123$"):
             Format.Enum(a, {1: 123})
+        with self.assertRaisesRegex(TypeError,
+                r"^Enum name must be a string or None, not 123$"):
+            Format.Enum(a, {}, name=123)
 
 
 class FormatStructTestCase(FHDLTestCase):
diff --git a/tests/test_lib_enum.py b/tests/test_lib_enum.py
index 806c31e..d2a0ffe 100644
--- a/tests/test_lib_enum.py
+++ b/tests/test_lib_enum.py
@@ -320,3 +320,13 @@ class EnumTestCase(FHDLTestCase):
             A = 1
             B = py_enum.member(2)
         self.assertIs(2, EnumB.B.value)
+
+    def test_format(self):
+        class EnumA(Enum, shape=unsigned(2)):
+            A = 0
+            B = 1
+
+        sig = Signal(EnumA)
+        self.assertRepr(EnumA.format(sig, ""), """
+        (format-enum (sig sig) 'EnumA' (0 'A') (1 'B'))
+        """)