usb-tools/amaranth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

lib.enum: add Enum wrappers that allow specifying shape.

Browse source 
See #756 and amaranth-lang/rfcs#3.
This commit is contained in:
Catherine 2023-02-20 22:58:38 +00:00
parent ef2e9fa809
commit 57612f1dce
10 changed files with 343 additions and 40 deletions
Download patch file Download diff file Expand all files Collapse all files

54
amaranth/hdl/ast.py
View file

@ -78,11 +78,34 @@ class Shape:
self.width = width
self.signed = signed
# The algorithm for inferring shape for standard Python enumerations is factored out so that
# `Shape.cast()` and Amaranth's `EnumMeta.as_shape()` can both use it.
@staticmethod
def _cast_plain_enum(obj):
signed = False
width = 0
for member in obj:
try:
member_shape = Const.cast(member.value).shape()
except TypeError as e:
raise TypeError("Only enumerations whose members have constant-castable "
"values can be used in Amaranth code")
if not signed and member_shape.signed:
signed = True
width = max(width + 1, member_shape.width)
elif signed and not member_shape.signed:
width = max(width, member_shape.width + 1)
else:
width = max(width, member_shape.width)
return Shape(width, signed)
@staticmethod
def cast(obj, *, src_loc_at=0):
while True:
if isinstance(obj, Shape):
return obj
elif isinstance(obj, ShapeCastable):
new_obj = obj.as_shape()
elif isinstance(obj, int):
return Shape(obj)
elif isinstance(obj, range):
@ -93,24 +116,9 @@ class Shape:
bits_for(obj.stop - obj.step, signed))
return Shape(width, signed)
elif isinstance(obj, type) and issubclass(obj, Enum):
signed = False
width = 0
for member in obj:
try:
member_shape = Const.cast(member.value).shape()
except TypeError as e:
raise TypeError("Only enumerations whose members have constant-castable "
"values can be used in Amaranth code")
if not signed and member_shape.signed:
signed = True
width = max(width + 1, member_shape.width)
elif signed and not member_shape.signed:
width = max(width, member_shape.width + 1)
else:
width = max(width, member_shape.width)
return Shape(width, signed)
elif isinstance(obj, ShapeCastable):
new_obj = obj.as_shape()
# For compatibility with third party enumerations, handle them as if they were
# defined as subclasses of lib.enum.Enum with no explicitly specified shape.
return Shape._cast_plain_enum(obj)
else:
raise TypeError("Object {!r} cannot be converted to an Amaranth shape".format(obj))
if new_obj is obj:
@ -866,9 +874,17 @@ class Cat(Value):
super().__init__(src_loc_at=src_loc_at)
self.parts = []
for index, arg in enumerate(flatten(args)):
if isinstance(arg, Enum) and (not isinstance(type(arg), ShapeCastable) or
not hasattr(arg, "_amaranth_shape_")):
warnings.warn("Argument #{} of Cat() is an enumerated value {!r} without "
"a defined shape used in bit vector context; define the enumeration "
"by inheriting from the class in amaranth.lib.enum and specifying "
"the 'shape=' keyword argument"
.format(index + 1, arg),
SyntaxWarning, stacklevel=2 + src_loc_at)
if isinstance(arg, int) and not isinstance(arg, Enum) and arg not in [0, 1]:
warnings.warn("Argument #{} of Cat() is a bare integer {} used in bit vector "
"context; consider specifying explicit width using C({}, {}) instead"
"context; specify the width explicitly using C({}, {})"
.format(index + 1, arg, arg, bits_for(arg)),
SyntaxWarning, stacklevel=2 + src_loc_at)
self.parts.append(Value.cast(arg))

108
amaranth/lib/enum.py Normal file
View file

@ -0,0 +1,108 @@
import enum as py_enum
import warnings
from ..hdl.ast import Shape, ShapeCastable, Const
from .._utils import bits_for
__all__ = py_enum.__all__
for member in py_enum.__all__:
globals()[member] = getattr(py_enum, member)
del member
class EnumMeta(ShapeCastable, py_enum.EnumMeta):
"""Subclass of the standard :class:`enum.EnumMeta` that implements the :class:`ShapeCastable`
protocol.
This metaclass provides the :meth:`as_shape` method, making its instances
:ref:`shape-castable <lang-shapecasting>`, and accepts a ``shape=`` keyword argument
to specify a shape explicitly. Other than this, it acts the same as the standard
:class:`enum.EnumMeta` class; if the ``shape=`` argument is not specified and
:meth:`as_shape` is never called, it places no restrictions on the enumeration class
or the values of its members.
"""
def __new__(metacls, name, bases, namespace, shape=None, **kwargs):
cls = py_enum.EnumMeta.__new__(metacls, name, bases, namespace, **kwargs)
if shape is not None:
# 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 = Shape.cast(shape)
for member in cls:
try:
Const.cast(member.value)
except TypeError as e:
raise TypeError("Value of enumeration member {!r} must be "
"a constant-castable expression"
.format(member)) from e
width = bits_for(member.value, shape.signed)
if member.value < 0 and not shape.signed:
warnings.warn(
message="Value of enumeration member {!r} is signed, but enumeration "
"shape is {!r}" # the repr will be `unsigned(X)`
.format(member, shape),
category=RuntimeWarning,
stacklevel=2)
elif width > shape.width:
warnings.warn(
message="Value of enumeration member {!r} will be truncated to "
"enumeration shape {!r}"
.format(member, shape),
category=RuntimeWarning,
stacklevel=2)
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
return cls
def as_shape(cls):
"""Cast this enumeration to a shape.
Returns
-------
:class:`Shape`
Explicitly provided shape. If not provided, returns the result of shape-casting
this class :ref:`as a standard Python enumeration <lang-shapeenum>`.
Raises
------
TypeError
If the enumeration has neither an explicitly provided shape nor any members.
"""
if hasattr(cls, "_amaranth_shape_"):
# Shape was provided explicitly; return it.
return cls._amaranth_shape_
elif cls.__members__:
# Shape was not provided explicitly, but enumeration has members; treat it
# the same way `Shape.cast` treats standard library enumerations, so that
# `amaranth.lib.enum.Enum` can be a drop-in replacement for `enum.Enum`.
return Shape._cast_plain_enum(cls)
else:
# Shape was not provided explicitly, and enumeration has no members.
# This is a base or mixin class that cannot be instantiated directly.
raise TypeError("Enumeration '{}.{}' does not have a defined shape"
.format(cls.__module__, cls.__qualname__))
class Enum(py_enum.Enum, metaclass=EnumMeta):
"""Subclass of the standard :class:`enum.Enum` that has :class:`EnumMeta` as
its metaclass."""
class IntEnum(py_enum.IntEnum, metaclass=EnumMeta):
"""Subclass of the standard :class:`enum.IntEnum` that has :class:`EnumMeta` as
its metaclass."""
class Flag(py_enum.Flag, metaclass=EnumMeta):
"""Subclass of the standard :class:`enum.Flag` that has :class:`EnumMeta` as
its metaclass."""
class IntFlag(py_enum.IntFlag, metaclass=EnumMeta):
"""Subclass of the standard :class:`enum.IntFlag` that has :class:`EnumMeta` as
its metaclass."""