summaryrefslogtreecommitdiffstats
path: root/aoc/decorators.py
blob: 98246e6e54bfaafd78484d4dbb0fbf80c1c19df3 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
# -*- coding: utf-8 -*-
from functools import wraps
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: AssignmentRun) -> AssignmentRun:
    @wraps(func)
    def wrapper(*args, **kwargs):
        items = list(func(*args, **kwargs))

        while True:
            for item in items:
                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