Return value of generator using itertools.chain

Question

I have the following code:

def my_generator():
  for i in range(10):
    yield i
  return 'I am a generator!'

def generator_wrapper():
  string = yield from my_generator()
  print(string)

list(generator_wrapper())

That outputs:

I am a generator!
Out[1]: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

I am trying to implement a decorator for my_generator(). I use itertools.chain because I would need the decorator to work on non-generators functions as well. Here is the code:

import itertools
import inspect

def _with_itertools(fn):
  def _impl():
    value = fn()
    if inspect.isgenerator(value):
      return itertools.chain([next(value)], value)
    else:
      return value
  return _impl


@_with_itertools
def my_generator():
  for i in range(10):
    yield i
  return 'I am a generator!'


def generator_wrapper():
  string = yield from my_generator()
  print(string)

list(generator_wrapper())

That outputs:

None
Out[2]: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

So using itertools.chain, I "lose" the return value of my generator.

Any ideas how to do it?

Answer 1

Reimplement chain in a way that preserves the StopIteration value:

def chain_with_StopIteration_value(*iterables):
    ret = None
    for it in iterables:
        ret = yield from it
    return ret

and use that instead of chain.