Can I memoize a Python generator?

I have a function called runquery that makes calls to a database and then yields the rows, one by one. I wrote a memoize de开发者_JAVA百科corator (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

---

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

---

from itertools import tee

sequence, memoized_sequence = tee (sequence, 2)

Done.

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

---

Yes. There's a decorator posted here. Take note that as the poster says, you lose some of the benefit of lazy evaluation.

def memoize(func):
    def inner(arg):
        if isinstance(arg, list):
            # Make arg immutable
            arg = tuple(arg)
        if arg in inner.cache:
            print "Using cache for %s" % repr(arg)
            for i in inner.cache[arg]:
                yield i
        else:
            print "Building new for %s" % repr(arg)
            temp = []
            for i in func(arg):
                temp.append(i)
                yield i
            inner.cache[arg] = temp
    inner.cache = {}
    return inner


@memoize
def gen(x):
    if not x:
        yield 0
        return

    for i in xrange(len(x)):
        for a in gen(x[i + 1:]):
            yield a + x[0]


print "Round 1"
for a in gen([2, 3, 4, 5]):
    print a

print
print "Round 2"
for a in gen([2, 3, 4, 5]):
    print a

---

Similar to the other answers, but simpler if you know that f is a generator:

def memoized_generator(f):
    cache = {}
    @functools.wraps(f)
    def wrapper(*args, **kwargs):
        k = args, frozenset(kwargs.items())
        it = cache[k] if k in cache else f(*args, **kwargs)
        cache[k], result = itertools.tee(it)
        return result
    return wrapper