Skip to content

Commit

Permalink
Don't insert spurious newlines in DummyTqdmFile.
Browse files Browse the repository at this point in the history
Only call `tqdm.write` (which inserts a newline) if
`DummyTqdmFile.write` is itself called with a string containing a
newline; otherwise buffer the rest of the string.

See changes in test_tqdm.py for the patterns this allows.
  • Loading branch information
anntzer committed Feb 15, 2021
1 parent 7f971d5 commit 285a5d2
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 4 deletions.
8 changes: 7 additions & 1 deletion tests/tests_tqdm.py
Expand Up @@ -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
Expand Down
15 changes: 12 additions & 3 deletions tqdm/contrib/__init__.py
Expand Up @@ -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):
Expand Down

0 comments on commit 285a5d2

Please sign in to comment.