Repository: amaranth-lang/amaranth
Pull request #1321 (2024-04-11 15:34:54)

lib.data: add `.format()` implementation.


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/data.py b/amaranth/lib/data.py
index 2cebc1b..ea35950 100644
--- a/amaranth/lib/data.py
+++ b/amaranth/lib/data.py
@@ -6,7 +6,7 @@ import operator
 
 from amaranth._utils import final
 from amaranth.hdl import *
-from amaranth.hdl._repr import *
+from amaranth.hdl._repr import Repr, FormatInt, FormatEnum
 from amaranth import hdl
 
 
@@ -248,6 +248,22 @@ class Layout(ShapeCastable, metaclass=ABCMeta):
         """
         return Const(self, raw)
 
+    def format(self, value, format_spec):
+        if format_spec != "":
+            raise ValueError(f"Format specifier {format_spec!r} is not supported for layouts")
+        value = Value.cast(value)
+        fields = {}
+        for key, field in self:
+            shape = Shape.cast(field.shape)
+            field_value = value[field.offset:field.offset+shape.width]
+            if isinstance(field.shape, ShapeCastable):
+                fields[str(key)] = field.shape.format(field.shape(field_value), "")
+            else:
+                if shape.signed:
+                    field_value = field_value.as_signed()
+                fields[str(key)] = Format("{}", field_value)
+        return Format.Struct(value, fields)
+
     def _value_repr(self, value):
         yield Repr(FormatInt(), value)
         for key, field in self:
@@ -539,6 +555,22 @@ class ArrayLayout(Layout):
     def __repr__(self):
         return f"ArrayLayout({self._elem_shape!r}, {self.length})"
 
+    def format(self, value, format_spec):
+        if format_spec != "":
+            raise ValueError(f"Format specifier {format_spec!r} is not supported for layouts")
+        value = Value.cast(value)
+        fields = []
+        shape = Shape.cast(self._elem_shape)
+        for index in range(self._length):
+            field_value = value[shape.width * index:shape.width * (index + 1)]
+            if isinstance(self._elem_shape, ShapeCastable):
+                fields.append(self._elem_shape.format(self._elem_shape(field_value), ""))
+            else:
+                if shape.signed:
+                    field_value = field_value.as_signed()
+                fields.append(Format("{}", field_value))
+        return Format.Array(value, fields)
+
 
 class FlexibleLayout(Layout):
     """Description of a flexible layout.
@@ -1163,6 +1195,9 @@ class _AggregateMeta(ShapeCastable, type):
     def from_bits(cls, bits):
         return cls.as_shape().from_bits(bits)
 
+    def format(cls, value, format_spec):
+        return cls.__layout.format(value, format_spec)
+
     def _value_repr(cls, value):
         return cls.__layout._value_repr(value)
 
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_data.py b/tests/test_lib_data.py
index 6618eea..3384935 100644
--- a/tests/test_lib_data.py
+++ b/tests/test_lib_data.py
@@ -75,7 +75,7 @@ class FieldTestCase(TestCase):
             data.Field(1, 0).offset = 1
 
 
-class StructLayoutTestCase(TestCase):
+class StructLayoutTestCase(FHDLTestCase):
     def test_construct(self):
         sl = data.StructLayout({
             "a": unsigned(1),
@@ -126,8 +126,21 @@ class StructLayoutTestCase(TestCase):
                 r"^Struct layout member shape must be a shape-castable object, not 1\.0$"):
             data.StructLayout({"a": 1.0})
 
+    def test_format(self):
+        sl = data.StructLayout({
+            "a": unsigned(1),
+            "b": signed(2),
+        })
+        sig = Signal(sl)
+        self.assertRepr(sl.format(sig, ""), """
+        (format-struct (sig sig)
+            ('a' (format '{}' (slice (sig sig) 0:1)))
+            ('b' (format '{}' (s (slice (sig sig) 1:3))))
+        )
+        """)
+
 
-class UnionLayoutTestCase(TestCase):
+class UnionLayoutTestCase(FHDLTestCase):
     def test_construct(self):
         ul = data.UnionLayout({
             "a": unsigned(1),
@@ -184,8 +197,21 @@ class UnionLayoutTestCase(TestCase):
                 r"\(specified: a, b\)$"):
             data.UnionLayout({"a": 1, "b": 2}).const(dict(a=1, b=2))
 
+    def test_format(self):
+        ul = data.UnionLayout({
+            "a": unsigned(1),
+            "b": 2
+        })
+        sig = Signal(ul)
+        self.assertRepr(ul.format(sig, ""), """
+        (format-struct (sig sig)
+            ('a' (format '{}' (slice (sig sig) 0:1)))
+            ('b' (format '{}' (slice (sig sig) 0:2)))
+        )
+        """)
 
-class ArrayLayoutTestCase(TestCase):
+
+class ArrayLayoutTestCase(FHDLTestCase):
     def test_construct(self):
         al = data.ArrayLayout(unsigned(2), 3)
         self.assertEqual(al.elem_shape, unsigned(2))
@@ -243,6 +269,48 @@ class ArrayLayoutTestCase(TestCase):
                 r"^Cannot index array layout with 'a'$"):
             al["a"]
 
+    def test_format(self):
+        al = data.ArrayLayout(unsigned(2), 3)
+        sig = Signal(al)
+        self.assertRepr(al.format(sig, ""), """
+        (format-array (sig sig)
+            (format '{}' (slice (sig sig) 0:2))
+            (format '{}' (slice (sig sig) 2:4))
+            (format '{}' (slice (sig sig) 4:6))
+        )
+        """)
+
+    def test_format_signed(self):
+        al = data.ArrayLayout(signed(2), 3)
+        sig = Signal(al)
+        self.assertRepr(al.format(sig, ""), """
+        (format-array (sig sig)
+            (format '{}' (s (slice (sig sig) 0:2)))
+            (format '{}' (s (slice (sig sig) 2:4)))
+            (format '{}' (s (slice (sig sig) 4:6)))
+        )
+        """)
+
+    def test_format_nested(self):
+        al = data.ArrayLayout(data.ArrayLayout(unsigned(2), 2), 3)
+        sig = Signal(al)
+        self.assertRepr(al.format(sig, ""), """
+        (format-array (sig sig)
+            (format-array (slice (sig sig) 0:4)
+                (format '{}' (slice (slice (sig sig) 0:4) 0:2))
+                (format '{}' (slice (slice (sig sig) 0:4) 2:4))
+            )
+            (format-array (slice (sig sig) 4:8)
+                (format '{}' (slice (slice (sig sig) 4:8) 0:2))
+                (format '{}' (slice (slice (sig sig) 4:8) 2:4))
+            )
+            (format-array (slice (sig sig) 8:12)
+                (format '{}' (slice (slice (sig sig) 8:12) 0:2))
+                (format '{}' (slice (slice (sig sig) 8:12) 2:4))
+            )
+        )
+        """)
+
 
 class FlexibleLayoutTestCase(TestCase):
     def test_construct(self):
@@ -1012,6 +1080,18 @@ class StructTestCase(FHDLTestCase):
         self.assertRepr(v.b.q.as_value(), "(slice (slice (sig v) 1:9) 4:8)")
         self.assertRepr(v.b.q.r, "(s (slice (slice (slice (sig v) 1:9) 4:8) 0:2))")
         self.assertRepr(v.b.q.s, "(s (slice (slice (slice (sig v) 1:9) 4:8) 2:4))")
+        self.assertRepr(S.format(v, ""), """
+        (format-struct (sig v)
+            ('a' (format '{}' (slice (sig v) 0:1)))
+            ('b' (format-struct (slice (sig v) 1:9)
+                ('p' (format '{}' (slice (slice (sig v) 1:9) 0:4)))
+                ('q' (format-struct (slice (slice (sig v) 1:9) 4:8)
+                    ('r' (format '{}' (s (slice (slice (slice (sig v) 1:9) 4:8) 0:2))))
+                    ('s' (format '{}' (s (slice (slice (slice (sig v) 1:9) 4:8) 2:4))))
+                ))
+            ))
+        )
+        """)
 
     def test_construct_init(self):
         class S(data.Struct):
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'))
+        """)