python tqdm与print的冲突及解决方法

问题

在开发fastnlp的过程中遇到以下问题：python中使用tqdm展示进度条，且同时打印一些输出时候会出现以下的情况：

import sys
import time

from tqdm import tqdm

for batch in tqdm(range(100), total=100, position=0, file=sys.stdout, desc="desc"):
    if batch % 5 == 0:
        print(batch)
        time.sleep(1)

原因

点击进入tqdm源码可以知道在__init__函数中有file的变量，若为None，则默认的输出流为sys.stderr。tqdm的进度条是采用file初始化的输出流来进行输出的。

python中自带的print函数默认的输出流为sys.stdout。而sys.stdout和sys.stderr的都是默认向屏幕输出，具体的区别是在默认情况下，stdout是行缓冲的，他的输出会放在一个buffer里面，只有到换行的时候，才会输出到屏幕。而stderr是无缓冲的，会直接输出。

故tqdm和print的打印冲突就是两者都向屏幕输出显示导致进度条和print打印的内容交杂。

解决方法

第一种方法

tqdm本身提供了write的方法，将print函数换为tqdm.write即可。

import sys
import time

from tqdm import tqdm

for batch in tqdm(range(100), total=100, position=0, file=sys.stdout, desc="desc"):
    if batch % 5 == 0:
        tqdm.write(str(batch))
        time.sleep(1)

不妨看看其解决方法：首先debug进入write函数内部，可以知道其跟print的区别是使用了with做了一些上下文处理，点击进入external_write_mode中，可以知道在write之前external_write_mode对已经展示的进度条进行clear操作，在write之后对复原clear的进度条。这就保证了进度条永远在用户打印的内容之下。

第二种方法

思路是将tqdm.write封装起来，对用户透明，用户还是能继续使用print函数。具体操作是将sys.stdout或者sys.stderr的write方法替换为tqdm.write方法，代码见下：

class DummyFile:
    def __init__(self, file):
        if file is None:
            file = sys.stderr
        self.file = file

    def write(self, x):
        if len(x.rstrip()) > 0:
            tqdm.write(x, file=self.file)

@contextlib.contextmanager
def redirect_stdout(file=None):
    if file is None:
        file = sys.stderr
    old_stdout = file
    sys.stdout = DummyFile(file)
    yield
    sys.stdout = old_stdout

for batch in tqdm(range(100), total=100, position=0, desc="desc"):
    if batch % 5 == 0:
        with redirect_stdout():
            print(batch)