TST: Use capture_stdout in testing

Author: gfyoung

pandas/tests/frame/test_repr_info.py

@@ -191,18 +191,16 @@ def test_latex_repr(self):
         # GH 12182
         self.assertIsNone(df._repr_latex_())
 
+    @tm.capture_stdout
     def test_info(self):
         io = StringIO()
         self.frame.info(buf=io)
         self.tsframe.info(buf=io)
 
         frame = DataFrame(np.random.randn(5, 3))
 
-        import sys
-        sys.stdout = StringIO()
         frame.info()
         frame.info(verbose=False)
-        sys.stdout = sys.__stdout__
 
     def test_info_wide(self):
         from pandas import set_option, reset_option

pandas/tests/io/parser/common.py

@@ -1254,6 +1254,7 @@ def test_regex_separator(self):
                              columns=['a', 'b', 'c'])
         tm.assert_frame_equal(result, expected)
 
+    @tm.capture_stdout
     def test_verbose_import(self):
         text = """a,b,c,d
 one,1,2,3
@@ -1265,22 +1266,19 @@ def test_verbose_import(self):
 one,1,2,3
 two,1,2,3"""
 
-        buf = StringIO()
-        sys.stdout = buf
-
-        try:  # engines are verbose in different ways
-            self.read_csv(StringIO(text), verbose=True)
-            if self.engine == 'c':
-                self.assertIn('Tokenization took:', buf.getvalue())
-                self.assertIn('Parser memory cleanup took:', buf.getvalue())
-            else:  # Python engine
-                self.assertEqual(buf.getvalue(),
-                                 'Filled 3 NA values in column a\n')
-        finally:
-            sys.stdout = sys.__stdout__
+        # Engines are verbose in different ways.
+        self.read_csv(StringIO(text), verbose=True)
+        output = sys.stdout.getvalue()
+
+        if self.engine == 'c':
+            self.assertIn('Tokenization took:', output)
+            self.assertIn('Parser memory cleanup took:', output)
+        else:  # Python engine
+            exp = 'Filled 3 NA values in column a\n'
+            self.assertEqual(output, exp)
 
-        buf = StringIO()
-        sys.stdout = buf
+        # Reset the stdout buffer.
+        sys.stdout = StringIO()
 
         text = """a,b,c,d
 one,1,2,3
@@ -1292,16 +1290,16 @@ def test_verbose_import(self):
 seven,1,2,3
 eight,1,2,3"""
 
-        try:  # engines are verbose in different ways
-            self.read_csv(StringIO(text), verbose=True, index_col=0)
-            if self.engine == 'c':
-                self.assertIn('Tokenization took:', buf.getvalue())
-                self.assertIn('Parser memory cleanup took:', buf.getvalue())
-            else:  # Python engine
-                self.assertEqual(buf.getvalue(),
-                                 'Filled 1 NA values in column a\n')
-        finally:
-            sys.stdout = sys.__stdout__
+        self.read_csv(StringIO(text), verbose=True, index_col=0)
+        output = sys.stdout.getvalue()
+
+        # Engines are verbose in different ways.
+        if self.engine == 'c':
+            self.assertIn('Tokenization took:', output)
+            self.assertIn('Parser memory cleanup took:', output)
+        else:  # Python engine
+            exp = 'Filled 1 NA values in column a\n'
+            self.assertEqual(output, exp)
 
     def test_iteration_open_handle(self):
         if PY3:

pandas/tests/io/parser/python_parser_only.py

@@ -8,7 +8,6 @@
 """
 
 import csv
-import sys
 import pytest
 
 import pandas.util.testing as tm
@@ -92,16 +91,9 @@ def test_BytesIO_input(self):
 
     def test_single_line(self):
         # see gh-6607: sniff separator
-
-        buf = StringIO()
-        sys.stdout = buf
-
-        try:
-            df = self.read_csv(StringIO('1,2'), names=['a', 'b'],
-                               header=None, sep=None)
-            tm.assert_frame_equal(DataFrame({'a': [1], 'b': [2]}), df)
-        finally:
-            sys.stdout = sys.__stdout__
+        df = self.read_csv(StringIO('1,2'), names=['a', 'b'],
+                           header=None, sep=None)
+        tm.assert_frame_equal(DataFrame({'a': [1], 'b': [2]}), df)
 
     def test_skipfooter(self):
         # see gh-6607

pandas/tests/io/test_sql.py

@@ -23,7 +23,6 @@
 import sqlite3
 import csv
 import os
-import sys
 
 import warnings
 import numpy as np
@@ -36,7 +35,7 @@
 from pandas import DataFrame, Series, Index, MultiIndex, isnull, concat
 from pandas import date_range, to_datetime, to_timedelta, Timestamp
 import pandas.compat as compat
-from pandas.compat import StringIO, range, lrange, string_types, PY36
+from pandas.compat import range, lrange, string_types, PY36
 from pandas.tseries.tools import format as date_format
 
 import pandas.io.sql as sql
@@ -2220,6 +2219,7 @@ def test_schema(self):
         cur = self.conn.cursor()
         cur.execute(create_sql)
 
+    @tm.capture_stdout
     def test_execute_fail(self):
         create_sql = """
         CREATE TABLE test
@@ -2236,14 +2236,11 @@ def test_execute_fail(self):
         sql.execute('INSERT INTO test VALUES("foo", "bar", 1.234)', self.conn)
         sql.execute('INSERT INTO test VALUES("foo", "baz", 2.567)', self.conn)
 
-        try:
-            sys.stdout = StringIO()
-            self.assertRaises(Exception, sql.execute,
-                              'INSERT INTO test VALUES("foo", "bar", 7)',
-                              self.conn)
-        finally:
-            sys.stdout = sys.__stdout__
+        self.assertRaises(Exception, sql.execute,
+                          'INSERT INTO test VALUES("foo", "bar", 7)',
+                          self.conn)
 
+    @tm.capture_stdout
     def test_execute_closed_connection(self):
         create_sql = """
         CREATE TABLE test
@@ -2259,12 +2256,10 @@ def test_execute_closed_connection(self):
 
         sql.execute('INSERT INTO test VALUES("foo", "bar", 1.234)', self.conn)
         self.conn.close()
-        try:
-            sys.stdout = StringIO()
-            self.assertRaises(Exception, tquery, "select * from test",
-                              con=self.conn)
-        finally:
-            sys.stdout = sys.__stdout__
+
+        self.assertRaises(Exception, tquery,
+                          "select * from test",
+                          con=self.conn)
 
         # Initialize connection again (needed for tearDown)
         self.setUp()
@@ -2534,6 +2529,7 @@ def test_schema(self):
         cur.execute(drop_sql)
         cur.execute(create_sql)
 
+    @tm.capture_stdout
     def test_execute_fail(self):
         _skip_if_no_pymysql()
         drop_sql = "DROP TABLE IF EXISTS test"
@@ -2553,14 +2549,11 @@ def test_execute_fail(self):
         sql.execute('INSERT INTO test VALUES("foo", "bar", 1.234)', self.conn)
         sql.execute('INSERT INTO test VALUES("foo", "baz", 2.567)', self.conn)
 
-        try:
-            sys.stdout = StringIO()
-            self.assertRaises(Exception, sql.execute,
-                              'INSERT INTO test VALUES("foo", "bar", 7)',
-                              self.conn)
-        finally:
-            sys.stdout = sys.__stdout__
+        self.assertRaises(Exception, sql.execute,
+                          'INSERT INTO test VALUES("foo", "bar", 7)',
+                          self.conn)
 
+    @tm.capture_stdout
     def test_execute_closed_connection(self):
         _skip_if_no_pymysql()
         drop_sql = "DROP TABLE IF EXISTS test"
@@ -2579,12 +2572,10 @@ def test_execute_closed_connection(self):
 
         sql.execute('INSERT INTO test VALUES("foo", "bar", 1.234)', self.conn)
         self.conn.close()
-        try:
-            sys.stdout = StringIO()
-            self.assertRaises(Exception, tquery, "select * from test",
-                              con=self.conn)
-        finally:
-            sys.stdout = sys.__stdout__
+
+        self.assertRaises(Exception, tquery,
+                          "select * from test",
+                          con=self.conn)
 
         # Initialize connection again (needed for tearDown)
         self.setUp()

pandas/util/decorators.py

@@ -1,7 +1,6 @@
-from pandas.compat import StringIO, callable, signature
+from pandas.compat import callable, signature
 from pandas._libs.lib import cache_readonly  # noqa
 import types
-import sys
 import warnings
 from textwrap import dedent
 from functools import wraps, update_wrapper
@@ -196,17 +195,6 @@ def indent(text, indents=1):
     return jointext.join(text.split('\n'))
 
 
-def suppress_stdout(f):
-    def wrapped(*args, **kwargs):
-        try:
-            sys.stdout = StringIO()
-            f(*args, **kwargs)
-        finally:
-            sys.stdout = sys.__stdout__
-
-    return wrapped
-
-
 def make_signature(func):
     """
     Returns a string repr of the arg list of a func call, with any defaults

pandas/util/testing.py

@@ -38,7 +38,7 @@
 from pandas.compat import (
     filter, map, zip, range, unichr, lrange, lmap, lzip, u, callable, Counter,
     raise_with_traceback, httplib, is_platform_windows, is_platform_32bit,
-    PY3
+    StringIO, PY3
 )
 
 from pandas.computation import expressions as expr
@@ -2790,3 +2790,20 @@ def setTZ(tz):
         yield
     finally:
         setTZ(orig_tz)
+
+
+def capture_stdout(f):
+    """
+    Capture stdout in a buffer so that it can be checked
+    (or suppressed) during testing.
+    """
+
+    @wraps(f)
+    def wrapped(*args, **kwargs):
+        try:
+            sys.stdout = StringIO()
+            f(*args, **kwargs)
+        finally:
+            sys.stdout = sys.__stdout__
+
+    return wrapped