基于Python标准库实现的tqdm时间估算优化方案咨询

我基于Python标准库实现了具备时间估算功能的基础版tqdm，代码如下：

import time
from datetime import datetime

def tqdm(iterable):
    steps = 0
    init_time = time()
    total_steps = len(iterable)
    last_update = init_time
    for i in iterable:
        yield i
        steps += 1
        if time() - last_update > 0.1:
            last_update = time()
            time_passed = last_update - init_time
            progress = steps / total_steps
            estimated_time = time_passed / progress - time_passed
            estimated_time = datetime.timedelta(seconds=int(estimated_time))
            print(f'\rEST: {estimated_time}', end='')

以官方tqdm模块为基准测试后，我发现自己实现的版本表现稍差，有时显示的估算时间比官方版偏长。请问有哪些方法可以优化我的代码，使其像官方tqdm一样健壮？

优化建议 #

1. 用滑动窗口平均计算速度，降低初始波动影响 #

当前实现依赖从启动到当前的总平均速度，初始阶段（前几步）的速度波动会导致估算严重失真。官方tqdm采用最近N个步骤的平均速度做估算，能有效平滑波动，提升准确性。

修改示例：

import time
from datetime import datetime
from collections import deque

def tqdm(iterable):
    steps = 0
    init_time = time()
    total_steps = len(iterable)
    last_update = init_time
    # 保存最近10个步骤的耗时，限制窗口大小
    step_times = deque(maxlen=10)
    prev_step_time = init_time

    for i in iterable:
        yield i
        steps += 1
        current_time = time()
        # 记录单步耗时
        step_times.append(current_time - prev_step_time)
        prev_step_time = current_time

        if current_time - last_update > 0.1:
            last_update = current_time
            if step_times:
                # 用滑动窗口的平均单步耗时计算剩余时间
                avg_step_time = sum(step_times) / len(step_times)
                remaining_steps = total_steps - steps
                estimated_time = avg_step_time * remaining_steps
            else:
                # 窗口未填满时回退到总平均
                time_passed = current_time - init_time
                progress = steps / total_steps
                estimated_time = (time_passed / progress) - time_passed
            
            # 避免估算时间为负数
            estimated_time = datetime.timedelta(seconds=int(max(estimated_time, 0)))
            print(f'\rEST: {estimated_time} | {steps}/{total_steps}', end='', flush=True)

2. 处理边界与异常情况 #

• 当steps极小（比如刚启动）时，progress趋近于0，会导致time_passed / progress出现极大值，需添加判断避免除以接近0的数；
• 兼容无长度的可迭代对象（如生成器），官方tqdm支持这种场景，此时可只显示已完成步数和实时速度，不做剩余时间估算：

def tqdm(iterable):
    steps = 0
    init_time = time()
    try:
        total_steps = len(iterable)
        has_total = True
    except TypeError:
        total_steps = None
        has_total = False
    
    # ... 其他逻辑 ...
    
    if has_total:
        # 计算并显示剩余时间
        print(f'\rEST: {estimated_time} | {steps}/{total_steps}', end='', flush=True)
    else:
        avg_speed = steps / (current_time - init_time)
        print(f'\rCompleted: {steps} | Speed: {avg_speed:.2f} steps/s', end='', flush=True)

3. 强制刷新输出缓冲区 #

默认print的end=''不会立即刷新终端，添加flush=True确保进度条实时更新，和官方版的即时反馈一致。

4. 增加中断清理逻辑 #

当迭代被中断（如Ctrl+C），官方tqdm会自动清理输出，避免后续命令与进度条重叠。可通过try...finally块实现：

def tqdm(iterable):
    try:
        # 原循环逻辑
        ...
    finally:
        # 结束时输出换行，还原终端格式
        print()

5. 减少冗余时间调用 #

避免在循环内频繁重复调用time()，比如只在需要更新进度或记录单步耗时的时候获取当前时间，降低性能开销。

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