使用itertools.tee复制嵌套迭代器(例如itertools.groupby)

8
我正在读取一个文件(同时进行一些昂贵的逻辑),我需要在不同的函数中多次迭代它,因此我真的希望只读取和解析该文件一次。
解析函数解析文件并返回一个itertools.groupby对象。
def parse_file():
    ...
    return itertools.groupby(lines, key=keyfunc)

我考虑做以下事情:

csv_file_content = read_csv_file()

file_content_1, file_content_2 = itertools.tee(csv_file_content, 2)

foo(file_content_1)
bar(file_content_2)

然而,itertools.tee 似乎只能“复制”外部迭代器,而内部(嵌套)迭代器仍然指向原始迭代器(因此在迭代通过 itertools.tee 返回的第一个迭代器后,它将被耗尽)。
独立的 MCVE:
from itertools import groupby, tee

li = [{'name': 'a', 'id': 1},
      {'name': 'a', 'id': 2},
      {'name': 'b', 'id': 3},
      {'name': 'b', 'id': 4},
      {'name': 'c', 'id': 5},
      {'name': 'c', 'id': 6}]

groupby_obj = groupby(li, key=lambda x:x['name'])
tee_obj1, tee_obj2 = tee(groupby_obj, 2)

print(id(tee_obj1))
for group, data in tee_obj1:
    print(group)
    print(id(data))
    for i in data:
        print(i)

print('----')

print(id(tee_obj2))
for group, data in tee_obj2:
    print(group)
    print(id(data))
    for i in data:
        print(i)

输出

2380054450440
a
2380053623136
{'name': 'a', 'id': 1}
{'name': 'a', 'id': 2}
b
2380030915976
{'name': 'b', 'id': 3}
{'name': 'b', 'id': 4}
c
2380054184344
{'name': 'c', 'id': 5}
{'name': 'c', 'id': 6}
----
2380064387336
a
2380053623136  # same ID as above
b
2380030915976  # same ID as above 
c
2380054184344  # same ID as above

如何高效地复制一个嵌套的迭代器?


1
但是如果你使用tee函数来复制内部迭代器,那么你不是会读取文件两次吗? - Dani Mesejo
您最好将所有内容硬编码到列表中。 - Jean-François Fabre
1
似乎即使在“tee”中,“grouped_object”也不能被使用两次。这个并行不起作用:“tee_obj1,tee_obj2 = groupby_obj,groupby_obj”。但我猜这会得到预期的结果:“tee_obj1,tee_obj2 = copy.deepcopy(groupby_obj),groupby_obj”。我猜.. - iGian
“如何递归地复制迭代器”无法得到恰当的答案(或者说没有解决方案),正如在此处所讨论的 https://dev59.com/SVgQ5IYBdhLWcg3w1njm,但在您的情况下,似乎`deepcopy`可以解决问题。 - Jean-François Fabre
1个回答

2

看起来 grouped_object (class 'itertools.groupby') 只能被消费一次,即使在 itertools.tee 中也是如此。 同时,相同的 grouped_object 的并行赋值不起作用:

tee_obj1, tee_obj2 = groupby_obj, groupby_obj

现在有效的是对 grouped_object 进行深层次的复制

tee_obj1, tee_obj2 = copy.deepcopy(groupby_obj), groupby_obj

2
似乎 grouped_objectct(类名为 itertools.groupby)只能被消耗一次,即使使用了 itertools.tee。但我不认为这是正确的,否则第二次就不会输出 a b c 了。虽然我希望使用比 deepcopy 更优雅的方法,但我还是接受了这个答案。 - DeepSpace
真实的情况是,每个groupby迭代返回的分组对象作为“值”时不能被tee。 - Jean-François Fabre

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接