summaryrefslogtreecommitdiffstats
path: root/aoc/decorators.py
diff options
context:
space:
mode:
Diffstat (limited to 'aoc/decorators.py')
-rw-r--r--aoc/decorators.py17
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 @@
1# -*- coding: utf-8 -*- 1# -*- coding: utf-8 -*-
2from functools import wraps 2from functools import wraps
3from typing import Iterator, TypeVar, Callable 3from typing import Iterator, TypeVar, Callable, List
4 4
5T = TypeVar("T") 5T = TypeVar("T")
6I = TypeVar("I")
6 7
8AssignmentRun = Callable[[Iterator[I], ...], Iterator[T]]
9AssignmentRunList = Callable[[List[I], ...], Iterator[T]]
7 10
8def infinite_generator(func: Callable[[...], Iterator[T]]): 11
12def infinite_generator(func: AssignmentRun) -> AssignmentRun:
9 @wraps(func) 13 @wraps(func)
10 def wrapper(*args, **kwargs): 14 def wrapper(*args, **kwargs):
11 items = list(func(*args, **kwargs)) 15 items = list(func(*args, **kwargs))
@@ -15,3 +19,12 @@ def infinite_generator(func: Callable[[...], Iterator[T]]):
15 yield item 19 yield item
16 20
17 return wrapper 21 return wrapper
22
23
24def list_input(func: AssignmentRun) -> AssignmentRunList:
25 @wraps(func)
26 def wrapper(self, input: Iterator, *args, **kwargs) -> T:
27
28 return func(self, list(input), *args, **kwargs)
29
30 return wrapper