From 9bb5a6270f7d4433b200c22ebc61984375910259 Mon Sep 17 00:00:00 2001 From: Daniel Sanchez Quiros Date: Wed, 6 Mar 2024 16:56:23 +0100 Subject: [PATCH] Implement chunks matrix transposed method --- da/common.py | 7 +++++-- da/test_common.py | 5 +++++ 2 files changed, 10 insertions(+), 2 deletions(-) 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)