TST: Use capture_stderr in testing

Author: gfyoung

pandas/tests/io/parser/common.py

@@ -1694,6 +1694,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'
@@ -1704,30 +1705,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()
+        self.assertEqual(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()
+        self.assertTrue('Skipping line 3' in val)
+        self.assertTrue('Skipping line 5' in val)

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,14 @@ 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()
+        self.assertTrue('Skipping line 4' in val)
+        self.assertTrue('Skipping line 6' in val)
 
     def test_header_not_enough_lines(self):
         data = ('skip this\n'

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()
+        self.assertEqual(output, '')
 
     def test_repr_name_iterable_indexable(self):
         s = Series([1, 2, 3], name=np.int64(3))

pandas/util/testing.py

@@ -2807,3 +2807,20 @@ def wrapped(*args, **kwargs):
             sys.stdout = sys.__stdout__
 
     return wrapped
+
+
+def capture_stderr(f):
+    """
+    Capture stderr in a buffer so that it can be checked
+    (or suppressed) during testing.
+    """
+
+    @wraps(f)
+    def wrapped(*args, **kwargs):
+        try:
+            sys.stderr = StringIO()
+            f(*args, **kwargs)
+        finally:
+            sys.stderr = sys.__stderr__
+
+    return wrapped