Python: More elegant solution for day 3

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]:
+def chunks(it: Iterable[T], n: int) -> Generator[list[T], None, None]:
-    if len(i) % 3 != 0:
+    if len(it) % n != 0:
+        raise ValueError()
+    if not 0 < n <= len(it):
         raise ValueError()
 
-    for idx in range(0, len(i), 3):
+    for idx in range(0, len(it), n):
-        yield i[idx], i[idx + 1], i[idx + 2]
+        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:

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