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]:
-    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: