captured iteration - yield and return

generator_util.py

from collections.abc import Iterator
from dataclasses import dataclass, field
from typing import Generic, Optional, TypeVar

T = TypeVar('T')
V = TypeVar('V')


@dataclass
class CapturedIteration(Generic[T, V]):
    generator: Iterator = field(repr=False)
    return_context: Optional[V] = None

    def __call__(self) -> Iterator[T]:
        while True:
            try:
                yield next(self.generator)
            except StopIteration as err:
                if err.args:
                    self.return_context = err.args[0]
                break


if __name__ == '__main__':
    # TODO: convert to unit test

    def fn():
        yield 'a'
        yield 'b'
        return 1

    dut: CapturedIteration[str, int] = CapturedIteration(generator=fn())
    for x in dut():
        print(x)
    print(dut.return_context)