我可以对Python生成器进行记忆吗？

Question

我有一个叫做 runquery 的函数，它会向数据库发起请求，然后一个一个地返回结果行。我写了一个记忆化装饰器（或者更准确地说，我是从这个StackOverflow问题上借来的），但是在后续的调用中，它只返回了空的序列，可能是因为生成器的值只能被返回一次。

我该如何修改这个记忆化装饰器，使它能适用于Python的生成器呢？我知道我需要在某个时候把结果存储在内存中，但我希望在装饰器内部处理这个问题，而不想修改原来的函数。

目前记忆化函数的代码是：

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

是的。这里有一个装饰器，链接在这里。需要注意的是，正如发帖者所说，你会失去一些延迟计算的好处。

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

Answer 2

from itertools import tee

sequence, memoized_sequence = tee (sequence, 2)

完成了。

对于生成器来说，这个事情更简单，因为标准库里有一个叫“tee”的方法！

Answer 3

我知道这个问题有点老了，不过对于那些想要完整解决方案的人来说，这里有一个，基于jsbueno的建议：

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