Can I memoize a Python generator?

Question

I have a function called runquery that makes calls to a database and then yields the rows, one by one. I wrote a memoize decorator (or more accurately, I just stole one from this stackoverflow question) but on subsequent calls it just yields an empty sequence, presumably because a generator's values can only be yielded once.

How could I modify the memoization decorator that works for Python generators? I realise I will need to store it in memory at some point but I'd like to handle this within the decorator and not modify the original function.

The current code of the memoization function is:

def memoized(f):     # Warning: Doesn't work if f yields values     cache={}     def ret(*args):         if args in cache:             return cache[args]         else:             answer=f(*args)             cache[args]=answer             return answer     return ret 

Answer 1

I realise this is somewhat of an old question, but for those who want a full solution: here's one, based on jsbueno's suggestion:

from itertools import tee from types import GeneratorType  Tee = tee([], 1)[0].__class__  def memoized(f):     cache={}     def ret(*args):         if args not in cache:             cache[args]=f(*args)         if isinstance(cache[args], (GeneratorType, Tee)):             # the original can't be used any more,             # so we need to change the cache as well             cache[args], r = tee(cache[args])             return r         return cache[args]     return ret 

Answer 2

from itertools import tee  sequence, memoized_sequence = tee (sequence, 2) 

Done.

It is easier for generators because the standard lib has this "tee" method!