Day 21 [WIP]main

1 files changed, 15 insertions, 2 deletions

diff --git a/aoc/decorators.py b/aoc/decorators.py
index a31993a..42603c2 100644
--- a/aoc/decorators.py
+++ b/aoc/decorators.py
@@ -1,11 +1,15 @@
 # -*- coding: utf-8 -*-
 from functools import wraps
-from typing import Iterator, TypeVar, Callable
+from typing import Iterator, TypeVar, Callable, List
 
 T = TypeVar("T")
+I = TypeVar("I")
 
+AssignmentRun = Callable[[Iterator[I], ...], Iterator[T]]
+AssignmentRunList = Callable[[List[I], ...], Iterator[T]]
 
-def infinite_generator(func: Callable[[...], Iterator[T]]):
+
+def infinite_generator(func: AssignmentRun) -> AssignmentRun:
     @wraps(func)
     def wrapper(*args, **kwargs):
         items = list(func(*args, **kwargs))
@@ -15,3 +19,12 @@ def infinite_generator(func: Callable[[...], Iterator[T]]):
                 yield item
 
     return wrapper
+
+
+def list_input(func: AssignmentRun) -> AssignmentRunList:
+    @wraps(func)
+    def wrapper(self, input: Iterator, *args, **kwargs) -> T:
+
+        return func(self, list(input), *args, **kwargs)
+
+    return wrapper