TST: Add decorators for redirecting stdout/err (#15952)

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,18 @@ def test_verbose_import(self):
 one,1,2,3
 two,1,2,3"""
 
-        buf = StringIO()
-        sys.stdout = buf
+        # Engines are verbose in different ways.
+        self.read_csv(StringIO(text), verbose=True)
+        output = sys.stdout.getvalue()
 
-        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__
+        if self.engine == 'c':
+            assert 'Tokenization took:' in output
+            assert 'Parser memory cleanup took:' in output
+        else:  # Python engine
+            assert output == 'Filled 3 NA values in column a\n'
 
-        buf = StringIO()
-        sys.stdout = buf
+        # Reset the stdout buffer.
+        sys.stdout = StringIO()
 
         text = """a,b,c,d
 one,1,2,3
@@ -1292,16 +1289,15 @@ 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':
+            assert 'Tokenization took:' in output
+            assert 'Parser memory cleanup took:' in output
+        else:  # Python engine
+            assert output == 'Filled 1 NA values in column a\n'
 
     def test_iteration_open_handle(self):
         if PY3:
@@ -1696,6 +1692,7 @@ class InvalidBuffer(object):
             with tm.assertRaisesRegexp(ValueError, msg):
                 self.read_csv(mock.Mock())
 
+    @tm.capture_stderr
     def test_skip_bad_lines(self):
         # see gh-15925
         data = 'a\n1\n1,2,3\n4\n5,6,7'
@@ -1706,30 +1703,24 @@ def test_skip_bad_lines(self):
         with tm.assertRaises(ParserError):
             self.read_csv(StringIO(data), error_bad_lines=True)
 
-        stderr = sys.stderr
         expected = DataFrame({'a': [1, 4]})
 
-        sys.stderr = StringIO()
-        try:
-            out = self.read_csv(StringIO(data),
-                                error_bad_lines=False,
-                                warn_bad_lines=False)
-            tm.assert_frame_equal(out, expected)
+        out = self.read_csv(StringIO(data),
+                            error_bad_lines=False,
+                            warn_bad_lines=False)
+        tm.assert_frame_equal(out, expected)
 
-            val = sys.stderr.getvalue()
-            self.assertEqual(val, '')
-        finally:
-            sys.stderr = stderr
+        val = sys.stderr.getvalue()
+        assert val == ''
 
+        # Reset the stderr buffer.
         sys.stderr = StringIO()
-        try:
-            out = self.read_csv(StringIO(data),
-                                error_bad_lines=False,
-                                warn_bad_lines=True)
-            tm.assert_frame_equal(out, expected)
 
-            val = sys.stderr.getvalue()
-            self.assertTrue('Skipping line 3' in val)
-            self.assertTrue('Skipping line 5' in val)
-        finally:
-            sys.stderr = stderr
+        out = self.read_csv(StringIO(data),
+                            error_bad_lines=False,
+                            warn_bad_lines=True)
+        tm.assert_frame_equal(out, expected)
+
+        val = sys.stderr.getvalue()
+        assert 'Skipping line 3' in val
+        assert 'Skipping line 5' in val

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/parser/test_textreader.py

@@ -142,6 +142,7 @@ def test_integer_thousands_alt(self):
         expected = DataFrame([123456, 12500])
         tm.assert_frame_equal(result, expected)
 
+    @tm.capture_stderr
     def test_skip_bad_lines(self):
         # too many lines, see #2430 for why
         data = ('a:b:c\n'
@@ -165,19 +166,15 @@ def test_skip_bad_lines(self):
                     2: ['c', 'f', 'i', 'n']}
         assert_array_dicts_equal(result, expected)
 
-        stderr = sys.stderr
-        sys.stderr = StringIO()
-        try:
-            reader = TextReader(StringIO(data), delimiter=':',
-                                header=None,
-                                error_bad_lines=False,
-                                warn_bad_lines=True)
-            reader.read()
-            val = sys.stderr.getvalue()
-            self.assertTrue('Skipping line 4' in val)
-            self.assertTrue('Skipping line 6' in val)
-        finally:
-            sys.stderr = stderr
+        reader = TextReader(StringIO(data), delimiter=':',
+                            header=None,
+                            error_bad_lines=False,
+                            warn_bad_lines=True)
+        reader.read()
+        val = sys.stderr.getvalue()
+
+        assert 'Skipping line 4' in val
+        assert 'Skipping line 6' in val
 
     def test_header_not_enough_lines(self):
         data = ('skip this\n'

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,10 @@ 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__
+        with pytest.raises(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 +2255,9 @@ 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__
+
+        with pytest.raises(Exception):
+            tquery("select * from test", con=self.conn)
 
         # Initialize connection again (needed for tearDown)
         self.setUp()
@@ -2534,6 +2527,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 +2547,10 @@ 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__
+        with pytest.raises(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 +2569,9 @@ 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__
+
+        with pytest.raises(Exception):
+            tquery("select * from test", con=self.conn)
 
         # Initialize connection again (needed for tearDown)
         self.setUp()

pandas/tests/plotting/test_frame.py

@@ -12,7 +12,7 @@
 from pandas import (Series, DataFrame, MultiIndex, PeriodIndex, date_range,
                     bdate_range)
 from pandas.types.api import is_list_like
-from pandas.compat import (range, lrange, StringIO, lmap, lzip, u, zip, PY3)
+from pandas.compat import range, lrange, lmap, lzip, u, zip, PY3
 from pandas.formats.printing import pprint_thing
 import pandas.util.testing as tm
 from pandas.util.testing import slow
@@ -1558,8 +1558,8 @@ def test_line_label_none(self):
         self.assertEqual(ax.get_legend().get_texts()[0].get_text(), 'None')
 
     @slow
+    @tm.capture_stdout
     def test_line_colors(self):
-        import sys
         from matplotlib import cm
 
         custom_colors = 'rgcby'
@@ -1568,16 +1568,13 @@ def test_line_colors(self):
         ax = df.plot(color=custom_colors)
         self._check_colors(ax.get_lines(), linecolors=custom_colors)
 
-        tmp = sys.stderr
-        sys.stderr = StringIO()
-        try:
-            tm.close()
-            ax2 = df.plot(colors=custom_colors)
-            lines2 = ax2.get_lines()
-            for l1, l2 in zip(ax.get_lines(), lines2):
-                self.assertEqual(l1.get_color(), l2.get_color())
-        finally:
-            sys.stderr = tmp
+        tm.close()
+
+        ax2 = df.plot(colors=custom_colors)
+        lines2 = ax2.get_lines()
+
+        for l1, l2 in zip(ax.get_lines(), lines2):
+            self.assertEqual(l1.get_color(), l2.get_color())
 
         tm.close()
 

pandas/tests/series/test_repr.py

@@ -3,13 +3,15 @@
 
 from datetime import datetime, timedelta
 
+import sys
+
 import numpy as np
 import pandas as pd
 
 from pandas import (Index, Series, DataFrame, date_range)
 from pandas.core.index import MultiIndex
 
-from pandas.compat import StringIO, lrange, range, u
+from pandas.compat import lrange, range, u
 from pandas import compat
 import pandas.util.testing as tm
 
@@ -112,20 +114,15 @@ def test_tidy_repr(self):
         a.name = 'title1'
         repr(a)  # should not raise exception
 
+    @tm.capture_stderr
     def test_repr_bool_fails(self):
         s = Series([DataFrame(np.random.randn(2, 2)) for i in range(5)])
 
-        import sys
-
-        buf = StringIO()
-        tmp = sys.stderr
-        sys.stderr = buf
-        try:
-            # it works (with no Cython exception barf)!
-            repr(s)
-        finally:
-            sys.stderr = tmp
-        self.assertEqual(buf.getvalue(), '')
+        # It works (with no Cython exception barf)!
+        repr(s)
+
+        output = sys.stderr.getvalue()
+        assert output == ''
 
     def test_repr_name_iterable_indexable(self):
         s = Series([1, 2, 3], name=np.int64(3))