gcgj-dify-1.7.0/api/libs/gen_utils.py

"""Utility functions for working with generators."""

from collections.abc import Callable, Generator
from typing import TypeVar

_YieldT = TypeVar("_YieldT")
_YieldR = TypeVar("_YieldR")


def map_(
    gen: Generator[_YieldT, None, None],
    mapper: Callable[[_YieldT], _YieldR],
) -> Generator[_YieldR, None, None]:
    for item in gen:
        yield mapper(item)


def filter_(
    gen: Generator[_YieldT, None, None],
    mapper: Callable[[_YieldT], bool],
) -> Generator[_YieldT, None, None]:
    for item in gen:
        if mapper(item):
            yield item


def wrap(
    gen: Generator[_YieldT, None, None],
    funcs: list[Callable[[Generator[_YieldT, None, None]], Generator[_YieldT, None, None]]],
) -> Generator[_YieldT, None, None]:
    for f in funcs:
        gen = f(gen)
    return gen