如何在python中保存迭代器的状态？

Posted 2021-04-19

我有一个非常大的迭代器。由于缺乏资源（网络，内存和时间），我不可能一步到位地执行我的程序。

所以我认为如果我将程序运行到迭代器中的第10000个元素然后保存其状态将会很好。下次我运行程序时它继续从迭代器中的第10001个元素开始。

这是我使用的代码：

import itertools
import requests

POSSIBLE_CHARS = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
URL = "URL?key={code}"

all_possible = itertools.combinations_with_replacement(POSSIBLE_CHARS, 29)
counter = itertools.count(start=1)

for c in all_possible:
    print("Try {}:".format(next(counter)), c)

    c_url = URL.format(code=c)
    resp = requests.get(c_url)
    if resp.status_code == 200:
        print("C found:", c)

        with open(c+".gif", 'b') as f:
            f.write(resp.content)

这个link展示了如何在单个执行中继续迭代器。但我想要的是停止程序并再次执行。

答案

所以我认为如果我将程序运行到迭代器中的第10000个元素然后保存其状态将会很好。下次我运行程序时它继续从迭代器中的第10001个元素开始。

你很幸运，因为itertools.combinations_with_replacement对象有API允许设置状态。这是由Python中的copy模块使用的，但没有理由你也不能挂钩。

例如，第10,001项将是这样的：

>>> all_possible = itertools.combinations_with_replacement(POSSIBLE_CHARS, 29)
>>> for i in range(10_000): 
...     next(all_possible) 
...
>>> "".join(next(all_possible))
'aaaaaaaaaaaaaaaaaaaaaaaaaafwI'

要在新实例中“快进”到此项，它将是：

>>> new_iterator = itertools.combinations_with_replacement(POSSIBLE_CHARS, 29)
>>> state = (0,)*26 + (5, 22, 33)
>>> new_iterator.__setstate__(state)
>>> "".join(next(new_iterator))
'aaaaaaaaaaaaaaaaaaaaaaaaaafwI'
>>> "".join(next(new_iterator))
'aaaaaaaaaaaaaaaaaaaaaaaaaafwJ'
>>> new_iterator.__setstate__(state)  # bonus: rewind iterator!
>>> "".join(next(new_iterator))
'aaaaaaaaaaaaaaaaaaaaaaaaaafwI'

要理解为什么组合10,001对应于长度为29的某个元组，如（0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0 ，0,0,0,0,0,0,0,0,22,33），我邀请您浏览CPython源to see how combinations objects allow themselves to be statefully copied以及pickle docs的相关部分。

如果您不需要针对错误或恶意构造的数据的安全性，那么您可以考虑简单地将这些迭代器腌制到文件而不是手动挂钩到__setstate__方法。