diff --git a/tests/tests_tqdm.py b/tests/tests_tqdm.py index dd33cc50..56033fef 100644 --- a/tests/tests_tqdm.py +++ b/tests/tests_tqdm.py @@ -1,6 +1,7 @@ # -*- coding: utf-8 -*- # Advice: use repr(our_file.read()) to print the full output of tqdm # (else '\r' will replace the previous lines and you'll see only the latest. +from __future__ import print_function import csv import os @@ -1722,8 +1723,14 @@ def test_file_redirection(): with closing(StringIO()) as our_file: # Redirect stdout to tqdm.write() with std_out_err_redirect_tqdm(tqdm_file=our_file): - for _ in trange(3): + with tqdm(total=3) as pbar: print("Such fun") + pbar.update(1) + print("Such", "fun") + pbar.update(1) + print("Such ", end="") + print("fun") + pbar.update(1) res = our_file.getvalue() assert res.count("Such fun\n") == 3 assert "0/3" in res diff --git a/tqdm/contrib/__init__.py b/tqdm/contrib/__init__.py index f9dfa171..8a6d2474 100644 --- a/tqdm/contrib/__init__.py +++ b/tqdm/contrib/__init__.py @@ -16,10 +16,28 @@ __all__ = ['tenumerate', 'tzip', 'tmap'] class DummyTqdmFile(ObjectWrapper): """Dummy file-like that will write to tqdm""" + def __init__(self, wrapped): + super(DummyTqdmFile, self).__init__(wrapped) + self._buf = [] + def write(self, x, nolock=False): - # Avoid print() second call (useless \n) - if len(x.rstrip()) > 0: - tqdm.write(x, file=self._wrapped, nolock=nolock) + nl = b"\n" if isinstance(x, bytes) else "\n" + pre, sep, post = x.rpartition(nl) + if sep: + blank = type(nl)() + tqdm.write(blank.join(self._buf + [pre, sep]), + end=blank, file=self._wrapped, nolock=nolock) + self._buf = [post] + else: + self._buf.append(x) + + def __del__(self): + if self._buf: + blank = type(self._buf[0])() + try: + tqdm.write(blank.join(self._buf), end=blank, file=self._wrapped) + except (OSError, ValueError): + pass def builtin_iterable(func):