diff --git a/amaranth/lib/data.py b/amaranth/lib/data.py index cd79c6e99..ec479b12f 100644 --- a/amaranth/lib/data.py +++ b/amaranth/lib/data.py @@ -796,17 +796,36 @@ def __getitem__(self, key): :exc:`TypeError` If :meth:`.ShapeCastable.__call__` does not return a value or a value-castable object. """ - if isinstance(key, slice): - raise TypeError( - "View cannot be indexed with a slice; did you mean to call `.as_value()` first?") if isinstance(self.__layout, ArrayLayout): - if not isinstance(key, (int, Value, ValueCastable)): + elem_width = Shape.cast(self.__layout.elem_shape).width + if isinstance(key, slice): + start, stop, stride = key.indices(self.__layout.length) + shape = ArrayLayout(self.__layout.elem_shape, len(range(start, stop, stride))) + if stride == 1: + value = self.__target[start * elem_width:stop * elem_width] + else: + value = Cat(self.__target[index * elem_width:(index + 1) * elem_width] + for index in range(start, stop, stride)) + elif isinstance(key, int): + if key not in range(-self.__layout.length, self.__layout.length): + raise IndexError(f"Index {key} is out of range for array layout of length " + f"{self.__layout.length}") + if key < 0: + key += self.__layout.length + shape = self.__layout.elem_shape + value = self.__target[key * elem_width:(key + 1) * elem_width] + elif isinstance(key, (int, Value, ValueCastable)): + shape = self.__layout.elem_shape + value = self.__target.word_select(key, elem_width) + else: raise TypeError( f"View with array layout may only be indexed with an integer or a value, " f"not {key!r}") - shape = self.__layout.elem_shape - value = self.__target.word_select(key, Shape.cast(self.__layout.elem_shape).width) else: + if isinstance(key, slice): + raise TypeError( + "Non-array view cannot be indexed with a slice; did you mean to call " + "`.as_value()` first?") if isinstance(key, (Value, ValueCastable)): raise TypeError( f"Only views with array layout, not {self.__layout!r}, may be indexed with " diff --git a/tests/test_lib_data.py b/tests/test_lib_data.py index 86a75c1da..a82fbf678 100644 --- a/tests/test_lib_data.py +++ b/tests/test_lib_data.py @@ -614,10 +614,52 @@ def test_getitem(self): self.assertEqual(v["q"].shape(), signed(1)) self.assertRepr(v["r"][0], "(slice (slice (sig v) 0:4) 0:2)") self.assertRepr(v["r"][1], "(slice (slice (sig v) 0:4) 2:4)") + self.assertRepr(v["r"][-2], "(slice (slice (sig v) 0:4) 0:2)") + self.assertRepr(v["r"][-1], "(slice (slice (sig v) 0:4) 2:4)") self.assertRepr(v["r"][i], "(part (slice (sig v) 0:4) (sig i) 2 2)") self.assertRepr(v["t"][0]["u"], "(slice (slice (slice (sig v) 0:4) 0:2) 0:1)") self.assertRepr(v["t"][1]["v"], "(slice (slice (slice (sig v) 0:4) 2:4) 1:2)") + def test_getitem_slice(self): + a = Signal(data.ArrayLayout(unsigned(2), 5)) + self.assertEqual(a[1:3].shape(), data.ArrayLayout(unsigned(2), 2)) + self.assertRepr(a[1:3].as_value(), "(slice (sig a) 2:6)") + self.assertRepr(a[2:].as_value(), "(slice (sig a) 4:10)") + self.assertRepr(a[:-2].as_value(), "(slice (sig a) 0:6)") + self.assertRepr(a[-1:].as_value(), "(slice (sig a) 8:10)") + self.assertRepr(a[::-1].as_value(), """ + (cat + (slice (sig a) 8:10) + (slice (sig a) 6:8) + (slice (sig a) 4:6) + (slice (sig a) 2:4) + (slice (sig a) 0:2) + ) + """) + self.assertRepr(a[::2].as_value(), """ + (cat + (slice (sig a) 0:2) + (slice (sig a) 4:6) + (slice (sig a) 8:10) + ) + """) + self.assertRepr(a[1::2].as_value(), """ + (cat + (slice (sig a) 2:4) + (slice (sig a) 6:8) + ) + """) + + def test_array_iter(self): + a = Signal(data.ArrayLayout(unsigned(2), 5)) + l = list(a) + self.assertEqual(len(l), 5) + self.assertRepr(l[0], "(slice (sig a) 0:2)") + self.assertRepr(l[1], "(slice (sig a) 2:4)") + self.assertRepr(l[2], "(slice (sig a) 4:6)") + self.assertRepr(l[3], "(slice (sig a) 6:8)") + self.assertRepr(l[4], "(slice (sig a) 8:10)") + def test_getitem_custom_call(self): class Reverser(ShapeCastable): def as_shape(self): @@ -674,9 +716,17 @@ def test_index_wrong_struct_dynamic(self): r"with a value$"): Signal(data.StructLayout({}))[Signal(1)] + def test_index_wrong_oob(self): + with self.assertRaisesRegex(IndexError, + r"^Index 2 is out of range for array layout of length 2$"): + Signal(data.ArrayLayout(unsigned(2), 2))[2] + with self.assertRaisesRegex(IndexError, + r"^Index -3 is out of range for array layout of length 2$"): + Signal(data.ArrayLayout(unsigned(2), 2))[-3] + def test_index_wrong_slice(self): with self.assertRaisesRegex(TypeError, - r"^View cannot be indexed with a slice; did you mean to call `.as_value\(\)` " + r"^Non-array view cannot be indexed with a slice; did you mean to call `.as_value\(\)` " r"first\?$"): Signal(data.StructLayout({}))[0:1]