lib.data: add .format() implementation.

Author: wanda-phi

amaranth/lib/data.py

@@ -6,7 +6,7 @@
 
 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,20 @@ def from_bits(self, raw):
         """
         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:
+                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 +553,20 @@ def size(self):
     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:
+                fields.append(Format("{}", field_value))
+        return Format.Array(value, fields)
+
 
 class FlexibleLayout(Layout):
     """Description of a flexible layout.
@@ -1163,6 +1191,9 @@ def const(cls, init):
     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)
 

tests/test_lib_data.py

@@ -75,7 +75,7 @@ def test_immutable(self):
             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 @@ def test_member_value_wrong(self):
                 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": 2
+        })
+        sig = Signal(sl)
+        self.assertRepr(sl.format(sig, ""), """
+        (format-struct (sig sig)
+            ('a' (format '{}' (slice (sig sig) 0:1)))
+            ('b' (format '{}' (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 @@ def test_const_two_members_wrong(self):
                 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,17 @@ def test_key_wrong_type(self):
                 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))
+        )
+        """)
+
 
 class FlexibleLayoutTestCase(TestCase):
     def test_construct(self):
@@ -1012,6 +1049,18 @@ class S(data.Struct):
         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 '{}' (slice (slice (slice (sig v) 1:9) 4:8) 0:2)))
+                    ('s' (format '{}' (slice (slice (slice (sig v) 1:9) 4:8) 2:4)))
+                ))
+            ))
+        )
+        """)
 
     def test_construct_init(self):
         class S(data.Struct):