From c0cdaa33ff377b1e6acf93afcf3618dd534073e1 2012-06-06 23:35:06 From: MinRK Date: 2012-06-06 23:35:06 Subject: [PATCH] move capture_output util from parallel tests to utils.io --- diff --git a/IPython/parallel/tests/clienttest.py b/IPython/parallel/tests/clienttest.py index 60017b6..478304d 100644 --- a/IPython/parallel/tests/clienttest.py +++ b/IPython/parallel/tests/clienttest.py @@ -100,36 +100,6 @@ def skip_without(*names): # Classes #------------------------------------------------------------------------------- -class CapturedIO(object): - """Simple object for containing captured stdout/err StringIO objects""" - - def __init__(self, stdout, stderr): - self.stdout_io = stdout - self.stderr_io = stderr - - @property - def stdout(self): - return self.stdout_io.getvalue() - - @property - def stderr(self): - return self.stderr_io.getvalue() - - -class capture_output(object): - """context manager for capturing stdout/err""" - - def __enter__(self): - self.sys_stdout = sys.stdout - self.sys_stderr = sys.stderr - stdout = sys.stdout = StringIO() - stderr = sys.stderr = StringIO() - return CapturedIO(stdout, stderr) - - def __exit__(self, exc_type, exc_value, traceback): - sys.stdout = self.sys_stdout - sys.stderr = self.sys_stderr - class ClusterTestCase(BaseZMQTestCase): diff --git a/IPython/parallel/tests/test_asyncresult.py b/IPython/parallel/tests/test_asyncresult.py index 9e0ba9c..e12f72c 100644 --- a/IPython/parallel/tests/test_asyncresult.py +++ b/IPython/parallel/tests/test_asyncresult.py @@ -18,11 +18,12 @@ Authors: import time -from IPython.parallel.error import TimeoutError +from IPython.utils.io import capture_output +from IPython.parallel.error import TimeoutError from IPython.parallel import error, Client from IPython.parallel.tests import add_engines -from .clienttest import ClusterTestCase, capture_output +from .clienttest import ClusterTestCase def setup(): add_engines(2, total=True) diff --git a/IPython/parallel/tests/test_magics.py b/IPython/parallel/tests/test_magics.py index 7279067..eeda8bd 100644 --- a/IPython/parallel/tests/test_magics.py +++ b/IPython/parallel/tests/test_magics.py @@ -25,6 +25,7 @@ from nose import SkipTest from IPython.testing import decorators as dec from IPython.testing.ipunittest import ParametricTestCase +from IPython.utils.io import capture_output from IPython import parallel as pmod from IPython.parallel import error @@ -33,7 +34,7 @@ from IPython.parallel.util import interactive from IPython.parallel.tests import add_engines -from .clienttest import ClusterTestCase, capture_output, generate_output +from .clienttest import ClusterTestCase, generate_output def setup(): add_engines(3, total=True) diff --git a/IPython/utils/io.py b/IPython/utils/io.py index 3869635..c3a25bc 100644 --- a/IPython/utils/io.py +++ b/IPython/utils/io.py @@ -17,6 +17,7 @@ from __future__ import print_function import os import sys import tempfile +from StringIO import StringIO #----------------------------------------------------------------------------- # Code @@ -321,3 +322,36 @@ def raw_print_err(*args, **kw): # Short aliases for quick debugging, do NOT use these in production code. rprint = raw_print rprinte = raw_print_err + + +class CapturedIO(object): + """Simple object for containing captured stdout/err StringIO objects""" + + def __init__(self, stdout, stderr): + self.stdout_io = stdout + self.stderr_io = stderr + + @property + def stdout(self): + return self.stdout_io.getvalue() + + @property + def stderr(self): + return self.stderr_io.getvalue() + + +class capture_output(object): + """context manager for capturing stdout/err""" + + def __enter__(self): + self.sys_stdout = sys.stdout + self.sys_stderr = sys.stderr + stdout = sys.stdout = StringIO() + stderr = sys.stderr = StringIO() + return CapturedIO(stdout, stderr) + + def __exit__(self, exc_type, exc_value, traceback): + sys.stdout = self.sys_stdout + sys.stderr = self.sys_stderr + + diff --git a/IPython/utils/tests/test_io.py b/IPython/utils/tests/test_io.py index eb49e40..71244d3 100644 --- a/IPython/utils/tests/test_io.py +++ b/IPython/utils/tests/test_io.py @@ -20,7 +20,7 @@ from subprocess import Popen, PIPE import nose.tools as nt from IPython.testing import decorators as dec -from IPython.utils.io import Tee +from IPython.utils.io import Tee, capture_output from IPython.utils.py3compat import doctest_refactor_print #----------------------------------------------------------------------------- @@ -73,3 +73,13 @@ def test_io_init(): # __class__ is a reference to the class object in Python 3, so we can't # just test for string equality. assert 'IPython.utils.io.IOStream' in classname, classname + +def test_capture_output(): + """capture_output() context works""" + + with capture_output() as io: + print 'hi, stdout' + print >> sys.stderr, 'hi, stderr' + + nt.assert_equals(io.stdout, 'hi, stdout\n') + nt.assert_equals(io.stderr, 'hi, stderr\n')