diff --git a/03/py/solution.py b/03/py/solution.py index 3ac9748..a05750e 100644 --- a/03/py/solution.py +++ b/03/py/solution.py @@ -6,12 +6,14 @@ from collections.abc import Iterable T = TypeVar("T") -def threes(i: Iterable[T]) -> Generator[tuple[T, T, T], None, None]: - if len(i) % 3 != 0: +def chunks(it: Iterable[T], n: int) -> Generator[list[T], None, None]: + if len(it) % n != 0: + raise ValueError() + if not 0 < n <= len(it): raise ValueError() - for idx in range(0, len(i), 3): - yield i[idx], i[idx + 1], i[idx + 2] + for idx in range(0, len(it), n): + yield it[idx : idx + n] def priority(s: str) -> int: @@ -44,7 +46,7 @@ def part1(data: str) -> int: def part2(data: str) -> int: result = 0 lines = data.splitlines() - for a, b, c in threes(lines): + for a, b, c in chunks(lines, 3): badge = set(a).intersection(set(b)).intersection(set(c)) value = 0 for ch in badge: diff --git a/03/py/test.py b/03/py/test.py index 3b0bb54..3542a4e 100644 --- a/03/py/test.py +++ b/03/py/test.py @@ -1,5 +1,5 @@ from typing import ContextManager, Optional -from solution import part1, part2, priority, threes +from solution import part1, part2, priority, chunks from contextlib import nullcontext as no_raise @@ -38,17 +38,17 @@ def test_priority(s: str, expected: int) -> None: @pytest.mark.parametrize( "line,expected,raises", [ - ("a\nb\nc", [("a", "b", "c")], no_raise()), + ("a\nb\nc", [["a", "b", "c"]], no_raise()), ("a\nb\nc\nd", None, pytest.raises(ValueError)), ], ) -def test_threes( +def test_chunks( line: str, expected: Optional[list[tuple[str, str, str]]], raises: ContextManager, ): with raises: - actual = list(threes(line.splitlines())) + actual = list(chunks(line.splitlines(), 3)) if expected: assert actual == expected