迭代器分页需求：取N个元素后继续获取同时间戳元素

需求概述 #

有一个按时间戳升序排序的迭代器，每个元素包含时间戳和业务数据。需要实现一个函数，完成以下逻辑：

1. 先从迭代器中取出前n个元素
2. 继续取出所有时间戳与最后一个取出元素时间戳相同的后续元素
3. 分页续取时，只需传入timestamp > 最后一次取出的时间戳的新迭代器即可，不能依赖存储特定时间戳的元素数量（因为当前时间戳可能有新元素插入）

示例数据 #

迭代器元素如下：

(1, "foo"),
(2, "bar"),
(2, "bar2"),
(3, "baz"),
(3, "baz2"),
(3, "baz3"),
(5, "qux")

调用效果 #

• take_n_and_while_timestamp_matches(1)：返回1个元素（时间戳1的元素）
• take_n_and_while_timestamp_matches(3)：返回3个元素（时间戳1、2的所有元素）
• take_n_and_while_timestamp_matches(2)：返回3个元素（时间戳1、2的所有元素）
• take_n_and_while_timestamp_matches(4)：返回6个元素（时间戳1、2、3的所有元素）
• take_n_and_while_timestamp_matches(100)：返回全部7个元素

实现方案 #

方案1：依赖more_itertools的简洁实现 #

利用peekable包装迭代器，方便查看下一个元素而不取出，避免取出不匹配元素后无法放回的问题：

from more_itertools import peekable

def take_n_and_while_timestamp_matches(iterator, n):
    it = peekable(iterator)
    result = []
    
    # 取出前n个元素
    for _ in range(n):
        if it:
            result.append(next(it))
        else:
            break
    
    if not result:
        return result
    
    last_ts = result[-1][0]
    
    # 继续取出所有时间戳匹配的后续元素
    while it and it.peek()[0] == last_ts:
        result.append(next(it))
    
    return result

方案2：无第三方依赖实现 #

如果不想依赖外部库，可以返回取出的结果和剩余迭代器，保证后续分页能继续使用：

def take_n_and_while_timestamp_matches(iterator, n):
    result = []
    remaining_buffer = []
    
    # 取出前n个元素
    for _ in range(n):
        try:
            item = next(iterator)
            result.append(item)
        except StopIteration:
            return result, iter(remaining_buffer)
    
    if not result:
        return result, iterator
    
    last_ts = result[-1][0]
    
    # 收集后续时间戳匹配的元素，遇到不匹配的则存入缓冲
    while True:
        try:
            item = next(iterator)
            if item[0] == last_ts:
                result.append(item)
            else:
                remaining_buffer.append(item)
                break
        except StopIteration:
            break
    
    # 生成包含缓冲元素和剩余迭代器的新迭代器
    def remaining_iterator():
        yield from remaining_buffer
        yield from iterator
    
    return result, remaining_iterator()

使用示例 #

# 构造示例迭代器
data = [
    (1, "foo"),
    (2, "bar"),
    (2, "bar2"),
    (3, "baz"),
    (3, "baz2"),
    (3, "baz3"),
    (5, "qux")
]
it = iter(data)

# 第一次取数
result, remaining_it = take_n_and_while_timestamp_matches(it, 2)
print(result)  # 输出: [(1, 'foo'), (2, 'bar'), (2, 'bar2')]

# 续取下一页
next_result, next_remaining = take_n_and_while_timestamp_matches(remaining_it, 3)
print(next_result)  # 输出: [(3, 'baz'), (3, 'baz2'), (3, 'baz3')]

内容的提问来源于stack exchange，提问作者Ákos Vandra-Meyer