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

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?

CodePudding user response：

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

CodePudding user response：

I guess we can extract it from the StopIteration exception like so:

def output_and_return(iterator):
  output = []
  try:
    while True:
      output.append(next(iterator))
  except StopIteration as e:
    return output, e.value

Not the cleanest code but quite a bit better than the class alternative.