diff --git a/tests/tests_tqdm.py b/tests/tests_tqdm.py index dd33cc502..414eebe71 100644 --- a/tests/tests_tqdm.py +++ b/tests/tests_tqdm.py @@ -1722,8 +1722,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 f9dfa1713..f0b80c2a1 100644 --- a/tqdm/contrib/__init__.py +++ b/tqdm/contrib/__init__.py @@ -16,10 +16,19 @@ 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 = "\n" if isinstance(x, str) else b"\n" + pre, sep, post = x.rpartition(nl) + if sep: + tqdm.write(type(nl)().join(self._buf) + pre, + file=self._wrapped, nolock=nolock) + self._buf = [post] + else: + self._buf.append(x) def builtin_iterable(func):