Python中如何基于时间戳合并两个迭代器并按规则输出？

我需要合并两个输出时间戳的迭代器，按时间戳从小到大的顺序输出元素：

• 当A迭代器的元素小于B时，输出A的元素
• 当两者元素相等时，输出包含两个元素的元组
• 当其中一个迭代器耗尽后，输出另一个迭代器剩余的所有元素

最小示例与预期 #

测试生成器 #

# 这些生成器的输出为“时间戳”
def gen_even():
    for x in range(0, 11, 2):
        yield x
        
def gen_odd():
    for x in sorted(list(range(1, 15, 2)) + [6]):
        yield x

预期合并结果 #

[0, 1, 2, 3, 4, 5, (6, 6), 7, 8, 9, 10, 11, 13]

尝试的代码及报错 #

我写了以下代码，但在其中一个迭代器耗尽时触发了StopIteration异常：

gen1 = gen_even()
gen2 = gen_odd()

def gen_both(gen1, gen2):
    first = next(gen1)
    second = next(gen2)
    
    while True:   
        if first < second:
            yield first
            first = next(gen1)
        elif first == second:  
            yield first, second
            first = next(gen1)
            second = next(gen2)            
        else:
            yield second
            second = next(gen2)


gen = gen_both(gen1, gen2)

for i in gen:
    print(i)

报错输出 #

0
1
2
3
4
5
(6, 6)
7
8
9
10
---------------------------------------------------------------------------
StopIteration                             Traceback (most recent call last)
Cell In[8], line 11, in gen_both(gen1, gen2)
     10     yield first
---> 11     first = next(gen1)
     12 elif first == second:  

StopIteration: 

The above exception was the direct cause of the following exception:

RuntimeError                              Traceback (most recent call last)
Cell In[8], line 21
     18             second = next(gen2)
     20 gen = gen_both(gen1, gen2)
---> 21 for i in gen:
     22     print(i)

RuntimeError: generator raised StopIteration

解决方案 #

问题核心是未处理迭代器耗尽的情况：Python 3.7+不允许生成器内部抛出StopIteration（会转为RuntimeError），必须主动捕获该异常并处理剩余元素。以下是几种可靠实现：

方法一：逐次捕获StopIteration（逻辑清晰） #

def gen_both(gen1, gen2):
    has_a = has_b = True
    val_a = val_b = None
    
    # 初始化获取第一个元素
    try:
        val_a = next(gen1)
    except StopIteration:
        has_a = False
    try:
        val_b = next(gen2)
    except StopIteration:
        has_b = False

    while has_a and has_b:
        if val_a < val_b:
            yield val_a
            try:
                val_a = next(gen1)
            except StopIteration:
                has_a = False
        elif val_a == val_b:
            yield (val_a, val_b)
            try:
                val_a = next(gen1)
            except StopIteration:
                has_a = False
            try:
                val_b = next(gen2)
            except StopIteration:
                has_b = False
        else:
            yield val_b
            try:
                val_b = next(gen2)
            except StopIteration:
                has_b = False
    
    # 输出剩余元素
    if has_a:
        yield val_a
        yield from gen1
    if has_b:
        yield val_b
        yield from gen2

方法二：循环内捕获异常并直接处理剩余元素 #

def gen_both(gen1, gen2):
    try:
        val1 = next(gen1)
        val2 = next(gen2)
    except StopIteration:
        # 处理其中一个迭代器初始为空的情况
        if 'val1' in locals():
            yield val1
            yield from gen1
        if 'val2' in locals():
            yield val2
            yield from gen2
        return

    while True:
        if val1 < val2:
            yield val1
            try:
                val1 = next(gen1)
            except StopIteration:
                # gen1耗尽，输出gen2剩余元素
                yield val2
                yield from gen2
                return
        elif val1 == val2:
            yield (val1, val2)
            try:
                val1 = next(gen1)
                val2 = next(gen2)
            except StopIteration:
                # 检查剩余元素
                if 'val1' in locals() and not gen1.gi_yieldfrom:
                    yield val1
                    yield from gen1
                if 'val2' in locals() and not gen2.gi_yieldfrom:
                    yield val2
                    yield from gen2
                return
        else:
            yield val2
            try:
                val2 = next(gen2)
            except StopIteration:
                # gen2耗尽，输出gen1剩余元素
                yield val1
                yield from gen1
                return

测试以上代码，都会输出预期序列：

0
1
2
3
4
5
(6, 6)
7
8
9
10
11
13

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