diff --git a/da/common.py b/da/common.py index d8ff0de..b918b5f 100644 --- a/da/common.py +++ b/da/common.py @@ -1,6 +1,6 @@ from dataclasses import dataclass from itertools import chain -from typing import List, Generator +from typing import List, Generator, Self from eth2spec.eip7594.mainnet import Bytes32 @@ -24,7 +24,10 @@ class Row(List[Bytes32]): class ChunksMatrix(List[Row]): @property def columns(self) -> Generator[List[Chunk], None, None]: - yield from map(list, zip(*self)) + yield from map(Row, zip(*self)) + + def transposed(self) -> Self: + return ChunksMatrix(self.columns) diff --git a/da/test_common.py b/da/test_common.py index de9aac5..26ea40b 100644 --- a/da/test_common.py +++ b/da/test_common.py @@ -10,3 +10,8 @@ class TestCommon(TestCase): expected = [[1, 4, 7], [2, 5, 8], [3, 6, 9]] for c1, c2 in zip(expected, matrix.columns): self.assertEqual(c1, c2) + + def test_chunks_matrix_transposed(self): + matrix = ChunksMatrix([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) + expected = ChunksMatrix([[1, 4, 7], [2, 5, 8], [3, 6, 9]]) + self.assertEqual(matrix.transposed(), expected)