Search code examples
pythongeneratornested-function

Python generator yielding from nested non-generator function


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.


Solution

  • 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)