我可以记住一个 Python 生成器吗？

Posted 2023-02-23

【中文标题】我可以记住一个 Python 生成器吗？

【英文标题】：Can I memoize a Python generator?

【发布时间】：2011-06-01 19:02:10

【问题描述】：

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

如何修改适用于 Python 生成器的 memoization 装饰器？我意识到我需要在某个时候将它存储在内存中，但我想在装饰器中处理它而不是修改原始函数。

memoization函数的当前代码是：

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

【参考方案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

【讨论】：

感谢您的插图！我花了很长时间才了解使用tee 的方式。 但是我认为检查实例时存在问题：您应该针对collections.Iterable 进行测试，因为针对types.GeneratorType 的测试仅适用于一次：在返回缓存的迭代器（iterator.tee 对象）时第三次调用该函数，缓存将返回一个耗尽的迭代器。

你是对的！然而，针对collections.Iterable 的测试也是错误的，因为列表和字符串等也是可迭代的。我将这些集合更改为 (GeneratorType, _tee)，因此它也适用于 tee 对象。

是的，这确实有点矫枉过正，但是，在 Python 2.7 下，itertools 中没有 _tee 对象，itertools.tee 函数返回 itertools.tee 对象。显然，这不是他们的课程，因此针对itertools.tee 进行测试毫无意义（而且也不起作用），因此我对collections.Iterable 提出了建议。你在哪个 python 版本下测试过你的代码？

嗯，好的。我针对 Python 3.3 进行了测试。我找到了一个适用于 3.3 和 2.7 的解决方案。我会马上更新我的答案。

哇，太好了 :) 谢谢你的努力！

【参考方案2】：

from itertools import tee

sequence, memoized_sequence = tee (sequence, 2)

完成。

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

【讨论】：

您能否编辑您的答案以显示如何将其与上面的 memoize 功能集成？我想在产生序列的函数上方输入类似@memoize_generator 的内容。

【参考方案3】：

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

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

【讨论】：

+1。但是，它不支持对记忆结果的交叉访问。

这很好，但还有改进的余地。使 *args 成为装饰器参数：比一个强制参数更有用。并将列表添加到缓存 before "for" - 这将为在第一次运行完成之前使用迭代器腾出空间。

【参考方案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