Python generator yielding from nested non-generator function

Question

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.

@KellyBundy, the actual code that I'm working on is more complex. It has two similar, but not identical loops doing something (similar between the loops, but not identical, so I cannot turn them into one loop), followed by this common code, plus one step between those loops with another copy of the common code. — Vedran Šego
– Vedran Šego, Commented Jan 21, 2024 at 2:37

Kelly Bundy · Accepted Answer · 2024-01-21 03:21:39Z

3

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)

edited Jan 21, 2024 at 3:21

answered Jan 21, 2024 at 2:40

Kelly Bundy

28k7 gold badges34 silver badges76 bronze badges

Sign up to request clarification or add additional context in comments.

3 Comments

Vedran Šego Over a year ago

Sounds about right. Thank you. I'll ponder on it more and see if it fits my use case or I need to expand the example above. Side note: the return type of add_one should be the same as that of f, right?

Kelly Bundy Over a year ago

@VedranŠego Right. I don't use type hints, forgot about that.

Vedran Šego Over a year ago

This works exactly as needed. Thank you.

Collectives™ on Stack Overflow

Python generator yielding from nested non-generator function

1 Answer 1

3 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

3 Comments

Your Answer

Sign up or log in

Post as a guest

Related