pythonで予め次のyieldを準備するgeneratorを作りたい

from threading import Thread, Event
from queue import Queue
import contextlib


@contextlib.contextmanager
def thread_generator(generator, max_queue_size, *args, **kwargs):
    """予め次にyieldする値を別スレッドで用意するgenerator

    マルチスレッド処理なので計算速度は向上しない。バッチをロードする
    ジェネレーターなど、ブロッキングI/Oが発生しうるものなどに対して有効

    Args:
        generator (function): 非同期で処理したいジェネレーター
        max_queue_size (int): 保存するキューの最大サイズ
        *args: generatorに引き渡す引数
        **kwargs: generatorに引き渡すキーワード引数

    Yields:
        第一引数で渡したgeneratorがyieldする値

    Examples:
        >>> import time
        >>> t_start = time.time()
        >>> def sample_generator(x, n):  # n回xを返すジェネレーター
        ...    for i in range(n):
        ...        time.sleep(1)
        ...        yield x
        >>> with thread_generator(sample_generator, 3, x=5, n=3) as gen:
        ...        for i in gen:
        ...            time.sleep(1)
        ...            print(i)
        5
        5
        5
        >>> t_end = time.time()
        >>> assert round(t_end - t_start) == 4  # 約4秒で実行できる

    Note:
        Examplesのように、1秒かかる処理を3回実行する場合最大4秒かかります。
        これは、サブスレッドの終了タイミングが、with句を抜け、かつ、
        その時点でスレッド内で実行されていた処理が｢終了した後｣であるからです。
    """
    queue = Queue(maxsize=max_queue_size)
    event = Event()

    def queuerunner(*args, **kwargs):
        """generatorが止まるまで値を取得してenqueueし続ける処理"""
        for i in generator(*args, **kwargs):
            if not event.is_set():
                queue.put(i)
            else:  # (finnalyで) eventフラグが立ったら終了する
                return

    thread = Thread(target=queuerunner, args=args, kwargs=kwargs, daemon=True)

    def queuegetter():
        """dequeueするだけのジェネレーター"""
        # runnerが実行されているかキューが残っている間はyieldし続ける
        while thread.is_alive() or not queue.empty():
            yield queue.get()

    thread.start()

    try:
        yield queuegetter()
    finally:
        # with句を抜けるタイミングで確実にスレッドを終わらせる処理
        event.set()
        thread.join()