fix: Make sure split_chunks doesn’t get 0 as chunk size

This commit is contained in:
Gergely Polonkai 2023-11-03 09:13:25 +01:00
parent 995f0dabed
commit c8b07ef913
2 changed files with 16 additions and 3 deletions

View File

@ -44,14 +44,17 @@ def inc_nonce(nonce: bytes) -> bytes:
return bnum
def split_chunks(seq: Sequence[T], n: int) -> Generator[Sequence[T], None, None]:
def split_chunks(seq: Sequence[T], chunk_size: int) -> Generator[Sequence[T], None, None]:
"""Split sequence in equal-sized chunks.
The last chunk is not padded."""
if chunk_size <= 0:
raise ValueError("chunk_size must be greater than zero")
while seq:
yield seq[:n]
seq = seq[n:]
yield seq[:chunk_size]
seq = seq[chunk_size:]
def long_to_bytes(n: int, blocksize: int = 0) -> bytes:

View File

@ -49,6 +49,16 @@ def test_split_chunks_is_generator() -> None:
assert isinstance(split_chunks([], 1), GeneratorType)
@pytest.mark.parametrize("size", (-123, -1, 0))
def test_nonpositive_chunk_size(size: int) -> None:
"""Test if split_chunks() with non-positive chunk sizes raise an error"""
with pytest.raises(ValueError) as ctx:
list(split_chunks(b"", size))
assert str(ctx.value) == "chunk_size must be greater than zero"
@pytest.mark.parametrize(
"in_,chunksize,out",
(