heapq.merge与自定义迭代器类兼容问题及解决方案咨询

问题复现 #

以下代码在使用heapq.merge合并两个自定义迭代器实例时，出现了额外的重复输出：

import heapq

class SimpleIterator:
    a_list = [0,1,2,3,4,5]

    def __iter__(self):
        self.list_iter = iter(self.a_list)
        return self

    def __next__(self):
        return next(self.list_iter)**2

print(list(SimpleIterator()))
print(list(heapq.merge(SimpleIterator(), SimpleIterator())))

预期输出：

[0, 1, 4, 9, 16, 25]
[0, 0, 1, 1, 4, 4, 9, 9, 16, 16, 25, 25]

实际输出：

[0, 1, 4, 9, 16, 25]
[0, 0, 1, 1, 4, 4, 9, 9, 16, 16, 25, 25, 0, 1, 4, 9, 16, 25]

疑问解答 #

1. 迭代器类实现是否有误？ #

是的，你的SimpleIterator类违反了Python迭代器协议的隐含预期：

• 该类同时作为可迭代对象和迭代器（__iter__返回自身），但每次调用__iter__都会重置内部的list_iter状态。
• 标准迭代器的__iter__方法应该返回自身且不修改迭代状态，而可迭代对象的__iter__应该返回一个全新的迭代器实例，保证每次迭代都是独立的。

当heapq.merge只剩一个迭代器时，会执行yield from next.__self__，这相当于对该迭代器对象再次调用iter()，触发了你的__iter__方法，重置了迭代状态，导致重复输出。

2. merge函数存在问题？ #

heapq.merge的实现逻辑没有问题，它是为了优化性能：当只剩一个输入迭代器时，直接通过yield from输出剩余元素，避免不必要的堆操作。问题出在你的自定义迭代器不符合迭代器协议的预期，而非merge本身的bug。

3. merge函数不应与迭代器类配合使用？ #

不是不能配合，只要你的迭代器符合标准协议即可：迭代器的__iter__返回自身且不重置状态，或者可迭代对象的__iter__返回全新的迭代器实例。

解决方案 #

方案1：分离可迭代对象与迭代器 #

将原类拆分为可迭代对象和独立的迭代器实现，确保每次调用iter()都生成全新的迭代器实例：

import heapq

class SimpleIterable:
    a_list = [0,1,2,3,4,5]

    def __iter__(self):
        # 返回独立的迭代器实例
        return _SimpleIterator(self.a_list)

class _SimpleIterator:
    def __init__(self, lst):
        self.list_iter = iter(lst)
    
    def __iter__(self):
        return self
    
    def __next__(self):
        return next(self.list_iter)**2

print(list(SimpleIterable()))
print(list(heapq.merge(SimpleIterable(), SimpleIterable())))

方案2：修改原类的__iter__逻辑（不推荐，仅单次迭代有效） #

调整__iter__，只在实例初始化时创建迭代器，避免每次调用iter()重置状态：

import heapq

class SimpleIterator:
    a_list = [0,1,2,3,4,5]

    def __init__(self):
        self.list_iter = None

    def __iter__(self):
        if self.list_iter is None:
            self.list_iter = iter(self.a_list)
        return self

    def __next__(self):
        val = next(self.list_iter)
        return val**2

# 注意：同一个实例只能迭代一次，复用会得到空结果
print(list(SimpleIterator()))
print(list(heapq.merge(SimpleIterator(), SimpleIterator())))

方案3：保留你的生成器包装方案 #

你当前使用的生成器包装之所以有效，是因为生成器对象是标准迭代器，它的__iter__返回自身且不会重置状态，yield from生成器时不会触发状态重置：

def generator_wrapper(obj: SimpleIterator):
    yield from obj

print(list(generator_wrapper(SimpleIterator())))
print(list(heapq.merge(generator_wrapper(SimpleIterator()), generator_wrapper(SimpleIterator())))) 

内容的提问来源于stack exchange，提问作者Greg