From 204fc43d8e58b72a3843c81ef442e6fe6327703a Mon Sep 17 00:00:00 2001
From: Casper da Costa-Luis <tqdm@caspersci.uk.to>
Date: Mon, 29 May 2017 01:19:46 +0100
Subject: [PATCH] add tests and redirection example

---
 examples/redirect_print.py | 60 ++++++++++++++++++++++++++++++++++++++
 tqdm/tests/tests_tqdm.py   | 46 +++++++++++++++++++++++++++++
 2 files changed, 106 insertions(+)
 create mode 100644 examples/redirect_print.py

diff --git a/examples/redirect_print.py b/examples/redirect_print.py
new file mode 100644
index 00000000..7ef4aae6
--- /dev/null
+++ b/examples/redirect_print.py
@@ -0,0 +1,60 @@
+"""Redirecting writing
+
+If using a library that can print messages to the console, editing the library
+by  replacing `print()` with `tqdm.write()` may not be desirable.
+In that case, redirecting `sys.stdout` to `tqdm.write()` is an option.
+
+To redirect `sys.stdout`, create a file-like class that will write
+any input string to `tqdm.write()`, and supply the arguments
+`file=sys.stdout, dynamic_ncols=True`.
+
+A reusable canonical example is given below:
+"""
+from __future__ import print_function
+from time import sleep
+import contextlib
+import sys
+from tqdm import tqdm
+
+
+class DummyTqdmFile(object):
+    """Dummy file-like that will write to tqdm"""
+    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 stdout_redirect_to_tqdm():
+    save_stdout = sys.stdout
+    try:
+        sys.stdout = DummyTqdmFile(sys.stdout)
+        yield save_stdout
+    # Relay exceptions
+    except Exception as exc:
+        raise exc
+    # Always restore sys.stdout if necessary
+    finally:
+        sys.stdout = save_stdout
+
+
+def blabla():
+    print("Foo blabla")
+
+
+# Redirect stdout to tqdm.write() (don't forget the `as save_stdout`)
+with stdout_redirect_to_tqdm() as save_stdout:
+    # tqdm call need to specify sys.stdout, not sys.stderr (default)
+    # and dynamic_ncols=True to autodetect console width
+    for _ in tqdm(range(3), file=save_stdout, dynamic_ncols=True):
+        blabla()
+        sleep(.5)
+
+# After the `with`, printing is restored
+print('Done!')
diff --git a/tqdm/tests/tests_tqdm.py b/tqdm/tests/tests_tqdm.py
index ae0341e3..eefeefa7 100644
--- a/tqdm/tests/tests_tqdm.py
+++ b/tqdm/tests/tests_tqdm.py
@@ -11,6 +11,7 @@ from nose import with_setup
 from nose.plugins.skip import SkipTest
 from nose.tools import assert_raises
 from time import sleep
+from contextlib import contextmanager
 
 from tqdm import tqdm
 from tqdm import trange
@@ -1549,3 +1550,48 @@ def test_postfix():
 
     out3 = out3[1:-1].split(', ')[3:]
     assert out3 == expected_order
+
+
+class DummyTqdmFile(object):
+    """Dummy file-like that will write to tqdm"""
+    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)
+
+
+@contextmanager
+def stdout_stderr_redirect_to_tqdm(tqdm_file=sys.stderr):
+    save_stdout = sys.stdout
+    save_stderr = sys.stderr
+    try:
+        sys.stdout = DummyTqdmFile(tqdm_file)
+        sys.stderr = DummyTqdmFile(tqdm_file)
+        yield save_stdout, save_stderr
+    # Relay exceptions
+    except Exception as exc:
+        raise exc
+    # Always restore sys.stdout if necessary
+    finally:
+        sys.stdout = save_stdout
+        sys.stderr = save_stderr
+
+
+@with_setup(pretest, posttest)
+def test_file_redirection():
+    """Test redirection of output"""
+    with closing(StringIO()) as our_file:
+        # Redirect stdout to tqdm.write()
+        with stdout_stderr_redirect_to_tqdm(tqdm_file=our_file):
+            # as (stdout, stderr)
+            for _ in trange(3):
+                print("Such fun")
+        res = our_file.getvalue()
+        assert res.count("Such fun\n") == 3
+        assert "0/3" in res
+        assert "3/3" in res