Merge pull request #15430 from sethtroisi/refguide_warning

MAINT: Use contextmanager in _run_doctests

1 files changed, 32 insertions, 46 deletions

diff --git a/tools/refguide_check.py b/tools/refguide_check.py
index a10ab87f3..fbc95f500 100644
--- a/tools/refguide_check.py
+++ b/tools/refguide_check.py
@@ -25,21 +25,23 @@ or in RST-based documentations::
 
     $ python refguide_check.py --rst docs
 """
-import sys
-import os
-import re
 import copy
-import inspect
-import warnings
 import doctest
-import tempfile
+import glob
+import inspect
 import io
-import docutils.core
-from docutils.parsers.rst import directives
+import os
+import re
 import shutil
-import glob
-from doctest import NORMALIZE_WHITESPACE, ELLIPSIS, IGNORE_EXCEPTION_DETAIL
+import sys
+import tempfile
+import warnings
+import docutils.core
 from argparse import ArgumentParser
+from contextlib import contextmanager, redirect_stderr
+from doctest import NORMALIZE_WHITESPACE, ELLIPSIS, IGNORE_EXCEPTION_DETAIL
+
+from docutils.parsers.rst import directives
 from pkg_resources import parse_version
 
 import sphinx
@@ -784,37 +786,27 @@ def _run_doctests(tests, full_name, verbose, doctest_warnings):
     runner = DTRunner(full_name, checker=Checker(), optionflags=flags,
                       verbose=verbose)
 
-    output = []
+    output = io.StringIO(newline='')
     success = True
-    def out(msg):
-        output.append(msg)
-
-    class MyStderr:
-        """
-        Redirect stderr to the current stdout
-        """
-        def write(self, msg):
-            if doctest_warnings:
-                sys.stdout.write(msg)
-            else:
-                out(msg)
 
-        # a flush method is required when a doctest uses multiprocessing
-        # multiprocessing/popen_fork.py flushes sys.stderr
-        def flush(self):
-            if doctest_warnings:
-                sys.stdout.flush()
+    # Redirect stderr to the stdout or output
+    tmp_stderr = sys.stdout if doctest_warnings else output
+
+    @contextmanager
+    def temp_cwd():
+        cwd = os.getcwd()
+        tmpdir = tempfile.mkdtemp()
+        try:
+            os.chdir(tmpdir)
+            yield tmpdir
+        finally:
+            os.chdir(cwd)
+            shutil.rmtree(tmpdir)
 
     # Run tests, trying to restore global state afterward
-    old_printoptions = np.get_printoptions()
-    old_errstate = np.seterr()
-    old_stderr = sys.stderr
     cwd = os.getcwd()
-    tmpdir = tempfile.mkdtemp()
-    sys.stderr = MyStderr()
-    try:
-        os.chdir(tmpdir)
-
+    with np.errstate(), np.printoptions(), temp_cwd() as tmpdir, \
+            redirect_stderr(tmp_stderr):
         # try to ensure random seed is NOT reproducible
         np.random.seed(None)
 
@@ -829,18 +821,12 @@ def _run_doctests(tests, full_name, verbose, doctest_warnings):
             # Process our options
             if any([SKIPBLOCK in ex.options for ex in t.examples]):
                 continue
-            fails, successes = runner.run(t, out=out, clear_globs=False)
+            fails, successes = runner.run(t, out=output.write, clear_globs=False)
             if fails > 0:
                 success = False
             ns = t.globs
-    finally:
-        sys.stderr = old_stderr
-        os.chdir(cwd)
-        shutil.rmtree(tmpdir)
-        np.set_printoptions(**old_printoptions)
-        np.seterr(**old_errstate)
 
-    return success, output
+    return success, output.read()
 
 
 def check_doctests(module, verbose, ns=None,
@@ -902,7 +888,7 @@ def check_doctests(module, verbose, ns=None,
         if dots:
             output_dot('.' if success else 'F')
 
-        results.append((full_name, success, "".join(output)))
+        results.append((full_name, success, output))
 
         if HAVE_MATPLOTLIB:
             import matplotlib.pyplot as plt
@@ -1022,7 +1008,7 @@ def check_doctests_testfile(fname, verbose, ns=None,
     if dots:
         output_dot('.' if success else 'F')
 
-    results.append((full_name, success, "".join(output)))
+    results.append((full_name, success, output))
 
     if HAVE_MATPLOTLIB:
         import matplotlib.pyplot as plt