Redirect print command in python script through tqdm.write()

Pierre Schroeder picture Pierre Schroeder · May 2, 2016 · Viewed 24.1k times · Source

I'm using tqdm in Python to display console-progressbars in our scripts. However, I have to call functions which print messages to the console as well and which I can't change. In general, writing to the console while displaying progress bars in the console messes up the display like so:

from time import sleep
from tqdm import tqdm

def blabla():
  print "Foo blabla"

for k in tqdm(range(3)):
  blabla()
  sleep(.5)

This creates the output:

0%|                                           | 0/3 [00:00<?, ?it/s]Foo
blabla
33%|###########6                       | 1/3 [00:00<00:01,  2.00it/s]Foo
blabla
67%|#######################3           | 2/3 [00:01<00:00,  2.00it/s]Foo
blabla
100%|###################################| 3/3 [00:01<00:00,  2.00it/s]

According to the documentation of tqdm the method tqdm.write() provides a means to write messages to the console without breaking the displayed progressbars. Thus, the right output is provided by this snippet:

from time import sleep
from tqdm import tqdm

def blabla():
  tqdm.write("Foo blabla")

for k in tqdm(range(3)):
  blabla()
  sleep(.5)

And looks like this:

Foo blabla
Foo blabla
Foo blabla
100%|###################################| 3/3 [00:01<00:00,  1.99it/s]

On the other hand, there is this solution which permits to silence those functions by quite elegantly redirecting sys.stdout into the void. This works perfectly well for silencing the functions.

Since I want to display the messages from these functions nonetheless without breaking the progress bars, I tried to merge both solutions into one by redirecting sys.stdout to tqdm.write() and, in turn, letting tqdm.write() write to the old sys.stdout. This results in the snippet:

from time import sleep

import contextlib
import sys

from tqdm import tqdm

class DummyFile(object):
  file = None
  def __init__(self, file):
    self.file = file

  def write(self, x):
    tqdm.write(x, file=self.file)

@contextlib.contextmanager
def nostdout():
    save_stdout = sys.stdout
    sys.stdout = DummyFile(save_stdout)
    yield
    sys.stdout = save_stdout

def blabla():
  print "Foo blabla"

for k in tqdm(range(3)):
  with nostdout():
    blabla()
    sleep(.5)

However, this actually creates an even more messed up output as before:

0%|                                           | 0/3 [00:00<?, ?it/s]Foo
blabla


33%|###########6                       | 1/3 [00:00<00:01,  2.00it/s]Foo
blabla


67%|#######################3           | 2/3 [00:01<00:00,  2.00it/s]Foo
blabla


100%|###################################| 3/3 [00:01<00:00,  2.00it/s]

FYI: calling tqdm.write(..., end="") inside DummyFile.write() creates the same result as the first output which is still messed up.

I can't understand why this wouldn't work, since tqdm.write() is supposed to manage clearing the progress bar before writing the message and then rewriting the progress bar.

What am I missing?

Answer

gaborous picture gaborous · May 15, 2016

Redirecting sys.stdout is always tricky, and it becomes a nightmare when two applications are twiddling with it at the same time.

Here the trick is that tqdm by default prints to sys.stderr, not sys.stdout. Normally, tqdm has an anti-mixup strategy for these two special channels, but since you are redirecting sys.stdout, tqdm gets confused because the file handle changes.

Thus, you just need to explicitly specify file=sys.stdout to tqdm and it will work:

from time import sleep

import contextlib
import sys

from tqdm import tqdm

class DummyFile(object):
  file = None
  def __init__(self, file):
    self.file = file

  def write(self, x):
    # Avoid print() second call (useless \n)
    if len(x.rstrip()) > 0:
        tqdm.write(x, file=self.file)

@contextlib.contextmanager
def nostdout():
    save_stdout = sys.stdout
    sys.stdout = DummyFile(sys.stdout)
    yield
    sys.stdout = save_stdout

def blabla():
  print("Foo blabla")

# tqdm call to sys.stdout must be done BEFORE stdout redirection
# and you need to specify sys.stdout, not sys.stderr (default)
for _ in tqdm(range(3), file=sys.stdout):
    with nostdout():
        blabla()
        sleep(.5)

print('Done!')

I also added a few more tricks to make the output nicer (eg, no useless \n when using print() without end='').

/EDIT: in fact it seems you can do the stdout redirection after starting tqdm, you just need to specify dynamic_ncols=True in tqdm.