From c5235b96af2fa0546ebb4369b18c3968430dc1e8 Mon Sep 17 00:00:00 2001 From: Guillaume Date: Mon, 5 Aug 2024 15:37:29 +0200 Subject: [PATCH] add `.split_every()` and `.transpose()` methods --- binapy/binapy.py | 33 +++++++++++++++++++++++++++++++++ tests/test_binapy.py | 7 +++++++ 2 files changed, 40 insertions(+) diff --git a/binapy/binapy.py b/binapy/binapy.py index 2bab570..ba93df5 100644 --- a/binapy/binapy.py +++ b/binapy/binapy.py @@ -6,6 +6,7 @@ import secrets from contextlib import suppress from functools import wraps +from itertools import zip_longest from typing import ( Any, Callable, @@ -339,6 +340,38 @@ def split_at(self, *pos: int) -> tuple[BinaPy, ...]: cut_at = split_at # for backward compatibility + def split_every(self, n: int, filler: bytes = b"") -> tuple[BinaPy, ...]: + """Split this BinaPy into chunks of size `n`. + + If filler is b'' (default), last chunk may be smaller than `n`. + Use a single char bytes instead for equally sized chunks. + + Args: + n: size of chunks + filler: char to use to fill the last chunk + + Returns: + a tuple of `len(self)/n` (rounded up) items. + + """ + div, mod = divmod(len(self), n) + if div == 0: + return (self + filler * (n - len(self)),) + if mod == 0: + return self.split_at(*range(n, len(self), n)) + chunks = self.split_at(*range(n, len(self), n)) + return *chunks[:-1], chunks[-1] + filler * (n - mod) + + def transpose(self, n: int, filler: bytes = b"") -> BinaPy: + """Reorder this BinaPy bytes from lines to columns.""" + + def iter_chunks() -> Iterator[bytes]: + lines = self.split_every(n, filler=filler) + for column in zip_longest(*lines): + yield bytes(c for c in column if c is not None) + + return BinaPy(b"".join(iter_chunks())) + extensions: ClassVar[dict[str, dict[str, Callable[..., Any]]]] = {} """Extension registry.""" diff --git a/tests/test_binapy.py b/tests/test_binapy.py index 09e2c1b..d705f95 100644 --- a/tests/test_binapy.py +++ b/tests/test_binapy.py @@ -62,6 +62,13 @@ def test_binapy() -> None: assert BinaPy.serialize_to("json", JSON).parse_from("json") == JSON + assert BinaPy("0123456789").split_at(3,5,8) == (b"012", b"34", b"567", b'89') + assert BinaPy("0123456789").split_every(4) == (b"0123", b"4567", b"89") + assert BinaPy("0123456789").split_every(12) == (b"0123456789",) + assert BinaPy("0123456789").split_every(4, filler=b' ') == (b"0123", b"4567", b"89 ") + + assert BinaPy("0123456789").transpose(3) == b"0369147258" + assert BinaPy("0123456789").transpose(3, filler=b' ') == b"0369147 258 " def test_slicing() -> None: bp = BinaPy("1234567890")