diff options
Diffstat (limited to 'aoc')
| -rw-r--r-- | aoc/decorators.py | 17 |
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 -*- |
| 2 | from functools import wraps | 2 | from functools import wraps |
| 3 | from typing import Iterator, TypeVar, Callable | 3 | from typing import Iterator, TypeVar, Callable, List |
| 4 | 4 | ||
| 5 | T = TypeVar("T") | 5 | T = TypeVar("T") |
| 6 | I = TypeVar("I") | ||
| 6 | 7 | ||
| 8 | AssignmentRun = Callable[[Iterator[I], ...], Iterator[T]] | ||
| 9 | AssignmentRunList = Callable[[List[I], ...], Iterator[T]] | ||
| 7 | 10 | ||
| 8 | def infinite_generator(func: Callable[[...], Iterator[T]]): | 11 | |
| 12 | def 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 | |||
| 24 | def 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 | ||
