lib.enum: use plain EnumMeta as metaclass when shape not used.

Author: wanda-phi

amaranth/lib/enum.py

@@ -83,16 +83,23 @@ def __new__(metacls, name, bases, namespace, shape=None, **kwargs):
                         category=SyntaxWarning,
                         stacklevel=2)
         # Actually instantiate the enumeration class.
-        cls = py_enum.EnumMeta.__new__(metacls, name, bases, namespace, **kwargs)
         if shape is not None:
+            cls = py_enum.EnumMeta.__new__(metacls, name, bases, namespace, **kwargs)
             # Shape is provided explicitly. Set the `_amaranth_shape_` attribute, and check that
             # the values of every member can be cast to the provided shape without truncation.
             cls._amaranth_shape_ = shape
         else:
             # Shape is not provided explicitly. Behave the same as a standard enumeration;
             # the lack of `_amaranth_shape_` attribute is used to emit a warning when such
             # an enumeration is used in a concatenation.
-            pass
+            bases = tuple(
+                py_enum.Enum if base is Enum else
+                py_enum.IntEnum if base is IntEnum else
+                py_enum.Flag if base is Flag else
+                py_enum.IntFlag if base is IntFlag else base
+                for base in bases
+            )
+            cls = py_enum.EnumMeta.__new__(py_enum.EnumMeta, name, bases, namespace, **kwargs)
         return cls
 
     def as_shape(cls):
@@ -144,21 +151,28 @@ def const(cls, init):
         return Const(member.value, cls.as_shape())
 
 
-class Enum(py_enum.Enum, metaclass=EnumMeta):
+class Enum(py_enum.Enum):
     """Subclass of the standard :class:`enum.Enum` that has :class:`EnumMeta` as
     its metaclass."""
 
 
-class IntEnum(py_enum.IntEnum, metaclass=EnumMeta):
+class IntEnum(py_enum.IntEnum):
     """Subclass of the standard :class:`enum.IntEnum` that has :class:`EnumMeta` as
     its metaclass."""
 
 
-class Flag(py_enum.Flag, metaclass=EnumMeta):
+class Flag(py_enum.Flag):
     """Subclass of the standard :class:`enum.Flag` that has :class:`EnumMeta` as
     its metaclass."""
 
 
-class IntFlag(py_enum.IntFlag, metaclass=EnumMeta):
+class IntFlag(py_enum.IntFlag):
     """Subclass of the standard :class:`enum.IntFlag` that has :class:`EnumMeta` as
     its metaclass."""
+
+# Fix up the metaclass after the fact: the metaclass __new__ requires these classes
+# to already be present, and also would not install itself on them due to lack of shape.
+Enum.__class__ = EnumMeta
+IntEnum.__class__ = EnumMeta
+Flag.__class__ = EnumMeta
+IntFlag.__class__ = EnumMeta

tests/test_hdl_ast.py

@@ -1077,7 +1077,7 @@ def const(self, init):
             Signal(CastableFromHex(), reset="01")
 
     def test_reset_shape_castable_enum_wrong(self):
-        class EnumA(AmaranthEnum):
+        class EnumA(AmaranthEnum, shape=1):
             X = 1
         with self.assertRaisesRegex(TypeError,
                 r"^Reset value must be a constant initializer of <enum 'EnumA'>$"):

tests/test_lib_enum.py

@@ -1,7 +1,7 @@
 import enum as py_enum
 
 from amaranth import *
-from amaranth.lib.enum import Enum
+from amaranth.lib.enum import Enum, EnumMeta
 
 from .utils import *
 
@@ -91,14 +91,13 @@ class EnumA(Enum, shape=unsigned(10)):
             A = 1
         self.assertRepr(Value.cast(EnumA.A), "(const 10'd1)")
 
-    def test_const_no_shape(self):
+    def test_no_shape(self):
         class EnumA(Enum):
             Z = 0
             A = 10
             B = 20
-        self.assertRepr(EnumA.const(None), "(const 5'd0)")
-        self.assertRepr(EnumA.const(10), "(const 5'd10)")
-        self.assertRepr(EnumA.const(EnumA.A), "(const 5'd10)")
+        self.assertNotIsInstance(EnumA, EnumMeta)
+        self.assertIsInstance(EnumA, py_enum.EnumMeta)
 
     def test_const_shape(self):
         class EnumA(Enum, shape=8):