我可以记住 Python 生成器吗？

Question

我有一个名为

runquery

的函数，它调用数据库，然后逐行生成行。我写了一个 memoize 装饰器（或者更准确地说，我只是从 this stackoverflow Question 中偷了一个），但在后续调用中它只产生一个空序列，大概是因为生成器的值只能产生一次。

如何修改适用于 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

我意识到这是一个有点老的问题，但对于那些想要完整解决方案的人来说：这是一个基于 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

Answer 2

from itertools import tee

sequence, memoized_sequence = tee (sequence, 2)

完成。

对于生成器来说更容易，因为标准库有这个“tee”方法！

Answer 3

是的。 这里发布了一个装饰器。请注意，正如海报所说，您会失去一些惰性评估的好处。

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 4

与其他答案类似，但如果您知道

f

是一个生成器，那就更简单了：

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

Answer 5

基于 jsbueno 和其他人（tee）的上述回答：

from functools import wraps

def cache_generator_func(f):
    g1, g2 = tee(f())
    @wraps(f)
    def wrapped():
        nonlocal g1, g2
        yield from g2
        g1, g2 = tee(g1)
    return wrapped

@cache_generator_func
def my_generator():
    "my generator"
    for i in range(5):
        sleep(1)
        yield i

%时间列表(my_generator())

CPU 时间：用户 2.29 ms，系统：96 μs，总计：2.39 ms 挂载时间：5.01 s [0,1,2,3,4]

%时间列表(my_generator())

CPU 时间：用户 5 μs，系统：0 ns，总计：5 μs 墙上时间：6.2 μs [0, 1, 2,3,4]

帮助（我的生成器）

模块 main 中函数 my_generator 的帮助：

我的生成器（） 我的发电机