diff --git a/tqdm/_tqdm.py b/tqdm/_tqdm.py index 38d8faf5..6b598d09 100644 --- a/tqdm/_tqdm.py +++ b/tqdm/_tqdm.py @@ -485,42 +485,45 @@ class tqdm(object): pass @classmethod - def write(cls, s, file=None, end="\n"): + def write(cls, s, file=None, end="\n", nolock=False): """ Print a message via tqdm (without overlap with bars) """ fp = file if file is not None else sys.stdout - with cls.external_write_mode(file=file): + with cls.external_write_mode(file=file, nolock=nolock): # Write the message fp.write(s) fp.write(end) @classmethod @contextmanager - def external_write_mode(cls, file=None): + def external_write_mode(cls, file=None, nolock=False): """ Disable tqdm within context and refresh tqdm when exits. Useful when writing to standard output stream """ fp = file if file is not None else sys.stdout - with cls._lock: - # Clear all bars - inst_cleared = [] - for inst in getattr(cls, '_instances', []): - # Clear instance if in the target output file - # or if write output + tqdm output are both either - # sys.stdout or sys.stderr (because both are mixed in terminal) - if inst.fp == fp or all( - f in (sys.stdout, sys.stderr) for f in (fp, inst.fp)): - inst.clear(nolock=True) - inst_cleared.append(inst) - yield - # Force refresh display of bars we cleared - for inst in inst_cleared: - # Avoid race conditions by checking that the instance started - if hasattr(inst, 'start_t'): # pragma: nocover - inst.refresh(nolock=True) + if not nolock: + cls._lock.acquire() + # Clear all bars + inst_cleared = [] + for inst in getattr(cls, '_instances', []): + # Clear instance if in the target output file + # or if write output + tqdm output are both either + # sys.stdout or sys.stderr (because both are mixed in terminal) + if inst.fp == fp or all( + f in (sys.stdout, sys.stderr) for f in (fp, inst.fp)): + inst.clear(nolock=True) + inst_cleared.append(inst) + yield + # Force refresh display of bars we cleared + for inst in inst_cleared: + # Avoid race conditions by checking that the instance started + if hasattr(inst, 'start_t'): # pragma: nocover + inst.refresh(nolock=True) + if not nolock: + cls._lock.release() # TODO: make list of all instances incl. absolutely positioned ones? @classmethod diff --git a/tqdm/tests/tests_tqdm.py b/tqdm/tests/tests_tqdm.py index 60d09d1d..8510086b 100644 --- a/tqdm/tests/tests_tqdm.py +++ b/tqdm/tests/tests_tqdm.py @@ -1601,7 +1601,7 @@ class DummyTqdmFile(object): def write(self, x): # Avoid print() second call (useless \n) if len(x.rstrip()) > 0: - tqdm.write(x, file=self.file) + tqdm.write(x, file=self.file, nolock=True) @contextmanager @@ -1636,6 +1636,20 @@ def test_file_redirection(): assert "3/3" in res +@with_setup(pretest, posttest) +def test_external_write(): + """Test external write mode""" + with closing(StringIO()) as our_file: + # Redirect stdout to tqdm.write() + for _ in trange(3, file=our_file): + with tqdm.external_write_mode(file=our_file): + our_file.write("Such fun\n") + res = our_file.getvalue() + assert res.count("Such fun\n") == 3 + assert "0/3" in res + assert "3/3" in res + + @with_setup(pretest, posttest) def test_unit_scale(): """Test numeric `unit_scale`"""