Search code examples
pythongenerator

Best way of getting both the yield'ed output and return'ed value of a generator in Python


Given a simple generator with a return value:

def my_generator():
    yield 1
    yield 2
    return 3

I'm looking for a simple function which returns the generated list and the return value.

>>> output_and_return(my_generator())
([1, 2], 3)

There doesn't seem to be any clean way of doing this. Inside another generator, you can use value = yield from my_generator(), which will get you the return value, but won't directly give you the output list itself.

The closest I can come up with is wrapping it around an iterator which captures the return value:

class Generator:
    def __init__(self, gen):
        self.gen = gen

    def __iter__(self):
        self.value = yield from self.gen
    
    def output_and_return(self):
        return list(self), self.value

Generator(my_generator()).output_and_return()

Which works, but it's anything but simple or clean. Does anyone know a simpler way of extracting the list of values and also the return value of a generator without wrapping it inside another class?


Solution

  • without wrapping it inside another class?

    Maybe with just a function instead?

    Version 1:

    def output_and_return(it):
        def with_result():
            yield (yield from it)
        *elements, result = with_result()
        return elements, result
    

    Version 2:

    def output_and_return(it):
        result = None
        def get_result():
            nonlocal result
            result = yield from it
        return list(get_result()), result