generator_decorator.py

import inspect
from typing import Callable, Generator, Awaitable, TypeVar, Union

YieldType = TypeVar("YieldType")
SendType = TypeVar("SendType")
ReturnType = TypeVar("ReturnType")

def friendly_generator_decorator(
    genfunction: Callable[..., Generator[YieldType, SendType, ReturnType]],
    handler: Callable[[YieldType], Union[Awaitable[SendType], SendType]],
) -> Callable[..., ReturnType]:
    """
    Convenient decorator for generator functions.
    Allows providing a handler that recieves/produces yields/sends from/to the generator.
    The result of the generator is then returned.
    handler can by sync or async.
    """
    if inspect.iscoroutinefunction(handler):
        async def wrapped_genfunction(*args, **kwargs):
            gen = genfunction(*args, **kwargs)
            try:
                yielded = next(gen)
                while True:
                    to_send = await handler(yielded)
                    yielded = gen.send(to_send)
            except StopIteration as e:
                return e.value

    else:
        def wrapped_genfunction(*args, **kwargs):
            gen = genfunction(*args, **kwargs)
            try:
                yielded = next(gen)
                while True:
                    to_send = handler(yielded)
                    yielded = gen.send(to_send)
            except StopIteration as e:
                return e.value

    return wrapped_genfunction


def g():
    yield 1
    return 2


async def handler(*args, **kwargs):
    print(args, kwargs)


if __name__ == "__main__":
    wrapped = friendly_generator_decorator(handler, g)
    print(wrapped())