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

后端 未结 4 666
抹茶落季
抹茶落季 2020-12-10 01:46

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

相关标签:
4条回答
  • 2020-12-10 02:10

    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.

    0 讨论(0)
  • 2020-12-10 02:12

    It might be the bad way, but I change built-in print function.

    import inspect
    import tqdm
    # store builtin print
    old_print = print
    def new_print(*args, **kwargs):
        # if tqdm.tqdm.write raises error, use builtin print
        try:
            tqdm.tqdm.write(*args, **kwargs)
        except:
            old_print(*args, ** kwargs)
    # globaly replace print with new_print
    inspect.builtins.print = new_print
    
    0 讨论(0)
  • 2020-12-10 02:19

    OP's solution is almost correct. The test in tqdm library that messes up your output is this one (https://github.com/tqdm/tqdm/blob/master/tqdm/_tqdm.py#L546-L549):

    if hasattr(inst, "start_t") and (inst.fp == fp or all(
               f in (sys.stdout, sys.stderr) for f in (fp, inst. 
        inst.clear(nolock=True)
        inst_cleared.append(inst)
    

    tqdm.write is testing the file you provided to see if there is risk of collision between the text to print and potential tqdm bars. In your case stdout and stderr get mixed in the terminal so there is a collision. To counter this, when the test passes, tqdm clears the bars, prints the text and draws the bars back after.

    Here, the test fp == sys.stdout fails because sys.stdout became DummyFile, and fp is the real sys.stdout, so the cleaning behavior is not enabled. A simple equality operator in DummyFile fixes everything.

    class DummyFile(object):
        def __init__(self, file):
            self.file = file
    
        def write(self, x):
            tqdm.write(x, end="", file=self.file)
    
        def __eq__(self, other):
            return other is self.file
    

    Also, since print passes a newline to sys.stdout (or not depending on the user choice), you don't want tqdm to add another one on its own, so it's better to set the option end='' than performing strip on the content.

    Advantages of this solution

    With gaborous's answer, tqdm(..., file=sys.stdout) pollutes your output stream with pieces of bar. By keeping file=sys.stdout (default), you keep your streams separated.
    With Conchylicultor and user493630's answers, you only patch print. However other systems such as logging directly stream to sys.stdout, so they don't go through tqdm.write.

    0 讨论(0)
  • 2020-12-10 02:22

    By mixing, user493630 and gaborous answers, I created this context manager which avoid having to use the file=sys.stdout parameter of tqdm.

    import inspect
    import contextlib
    import tqdm
    
    @contextlib.contextmanager
    def redirect_to_tqdm():
        # Store builtin print
        old_print = print
        def new_print(*args, **kwargs):
            # If tqdm.tqdm.write raises error, use builtin print
            try:
                tqdm.tqdm.write(*args, **kwargs)
            except:
                old_print(*args, ** kwargs)
    
        try:
            # Globaly replace print with new_print
            inspect.builtins.print = new_print
            yield
        finally:
            inspect.builtins.print = old_print
    

    To use it, simply:

    for i in tqdm.tqdm(range(100)):
        with redirect_to_tqdm():
            time.sleep(.1)
            print(i)
    

    To simplify even further, it is possible to wrap the code in a new function:

    def tqdm_redirect(*args, **kwargs):
        with redirect_to_tqdm():
            for x in tqdm.tqdm(*args, **kwargs):
                yield x
    
    for i in tqdm_redirect(range(20)):
        time.sleep(.1)
        print(i)
    
    0 讨论(0)
提交回复
热议问题