Python generator yielding from nested non-generator function

50 views Asked by At

This is a dumb example based on a more complex thing that I want to do:

from typing import Generator


def f() -> Generator[list[int], None, None]:
    result = list()
    result.append(1)
    if len(result) == 2:
        yield result
        result = list()
    result.append(2)
    if len(result) == 2:
        yield result
        result = list()
    result.append(3)
    if len(result) == 2:
        yield result
        result = list()
    result.append(4)
    if len(result) == 2:
        yield result
        result = list()


print(list(f()))

The point here is that this bit is copied multiple times:

    if len(result) == 2:
        yield result
        result = list()

Normally, I'd change it into something like this:

from typing import Generator


def f() -> Generator[list[int], None, None]:
    def add_one(value: int) -> None:
        nonlocal result
        result.append(value)
        if len(result) == 2:
            nonlocal_yield result
            result = list()

    result = list()
    add_one(1)
    add_one(2)
    add_one(3)
    add_one(4)


print(list(f()))

Obviously, nonlocal_yield is not a thing. Is there an elegant way to achieve this?

I know that I can just create the full list of results, i.e., [[1, 2], [3, 4]], and then either return it or yield individual 2-element sublists. Something like this:

from typing import Generator


def f() -> list[list[int]]:
    def add_one(value: int) -> None:
        nonlocal current
        current.append(value)
        if len(current) == 2:
            result.append(current)
            current = list()

    result = list()
    current = list()
    add_one(1)
    add_one(2)
    add_one(3)
    add_one(4)
    return result


print(list(f()))

However, this beats the purpose of a generator. I'll go for it in absence of a better solution, but I'm curious if there is a "pure" generator way to do it.

1

There are 1 answers

3
Kelly Bundy On BEST ANSWER

One possibility:

def f() -> Generator[list[int], None, None]:
    def add_one(value: int) -> Generator[list[int], None, None]:
        nonlocal result
        result.append(value)
        if len(result) == 2:
            yield result
            result = list()

    result = list()
    yield from add_one(1)
    yield from add_one(2)
    yield from add_one(3)
    yield from add_one(4)