move capture_output util from parallel tests to utils.io

IPython/parallel/tests/clienttest.py

@@ -1,213 +1,183 b''
 """base class for parallel client tests
 
 Authors:
 
 * Min RK
 """
 
 #-------------------------------------------------------------------------------
 #  Copyright (C) 2011  The IPython Development Team
 #
 #  Distributed under the terms of the BSD License.  The full license is in
 #  the file COPYING, distributed as part of this software.
 #-------------------------------------------------------------------------------
 
 import sys
 import tempfile
 import time
 from StringIO import StringIO
 
 from nose import SkipTest
 
 import zmq
 from zmq.tests import BaseZMQTestCase
 
 from IPython.external.decorator import decorator
 
 from IPython.parallel import error
 from IPython.parallel import Client
 
 from IPython.parallel.tests import launchers, add_engines
 
 # simple tasks for use in apply tests
 
 def segfault():
     """this will segfault"""
     import ctypes
     ctypes.memset(-1,0,1)
 
 def crash():
     """from stdlib crashers in the test suite"""
     import types
     if sys.platform.startswith('win'):
         import ctypes
         ctypes.windll.kernel32.SetErrorMode(0x0002);
     args = [ 0, 0, 0, 0, b'\x04\x71\x00\x00', (), (), (), '', '', 1, b'']
     if sys.version_info[0] >= 3:
         # Python3 adds 'kwonlyargcount' as the second argument to Code
         args.insert(1, 0)
         
     co = types.CodeType(*args)
     exec(co)
 
 def wait(n):
     """sleep for a time"""
     import time
     time.sleep(n)
     return n
 
 def raiser(eclass):
     """raise an exception"""
     raise eclass()
 
 def generate_output():
     """function for testing output
     
     publishes two outputs of each type, and returns
     a rich displayable object.
     """
     
     import sys
     from IPython.core.display import display, HTML, Math
     
     print "stdout"
     print >> sys.stderr, "stderr"
     
     display(HTML("<b>HTML</b>"))
     
     print "stdout2"
     print >> sys.stderr, "stderr2"
     
     display(Math(r"\alpha=\beta"))
     
     return Math("42")
 
 # test decorator for skipping tests when libraries are unavailable
 def skip_without(*names):
     """skip a test if some names are not importable"""
     @decorator
     def skip_without_names(f, *args, **kwargs):
         """decorator to skip tests in the absence of numpy."""
         for name in names:
             try:
                 __import__(name)
             except ImportError:
                 raise SkipTest
         return f(*args, **kwargs)
     return 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):
     
     def add_engines(self, n=1, block=True):
         """add multiple engines to our cluster"""
         self.engines.extend(add_engines(n))
         if block:
             self.wait_on_engines()
 
     def minimum_engines(self, n=1, block=True):
         """add engines until there are at least n connected"""
         self.engines.extend(add_engines(n, total=True))
         if block:
             self.wait_on_engines()
             
     
     def wait_on_engines(self, timeout=5):
         """wait for our engines to connect."""
         n = len(self.engines)+self.base_engine_count
         tic = time.time()
         while time.time()-tic < timeout and len(self.client.ids) < n:
             time.sleep(0.1)
         
         assert not len(self.client.ids) < n, "waiting for engines timed out"
     
     def connect_client(self):
         """connect a client with my Context, and track its sockets for cleanup"""
         c = Client(profile='iptest', context=self.context)
         for name in filter(lambda n:n.endswith('socket'), dir(c)):
             s = getattr(c, name)
             s.setsockopt(zmq.LINGER, 0)
             self.sockets.append(s)
         return c
     
     def assertRaisesRemote(self, etype, f, *args, **kwargs):
         try:
             try:
                 f(*args, **kwargs)
             except error.CompositeError as e:
                 e.raise_exception()
         except error.RemoteError as e:
             self.assertEquals(etype.__name__, e.ename, "Should have raised %r, but raised %r"%(etype.__name__, e.ename))
         else:
             self.fail("should have raised a RemoteError")
             
     def _wait_for(self, f, timeout=10):
         """wait for a condition"""
         tic = time.time()
         while time.time() <= tic + timeout:
             if f():
                 return
             time.sleep(0.1)
             self.client.spin()
         if not f():
             print "Warning: Awaited condition never arrived"
     
     def setUp(self):
         BaseZMQTestCase.setUp(self)
         self.client = self.connect_client()
         # start every test with clean engine namespaces:
         self.client.clear(block=True)
         self.base_engine_count=len(self.client.ids)
         self.engines=[]
     
     def tearDown(self):
         # self.client.clear(block=True)
         # close fds:
         for e in filter(lambda e: e.poll() is not None, launchers):
             launchers.remove(e)
         
         # allow flushing of incoming messages to prevent crash on socket close
         self.client.wait(timeout=2)
         # time.sleep(2)
         self.client.spin()
         self.client.close()
         BaseZMQTestCase.tearDown(self)
         # this will be redundant when pyzmq merges PR #88
         # self.context.term()
         # print tempfile.TemporaryFile().fileno(),
         # sys.stdout.flush()
         
  No newline at end of file

IPython/parallel/tests/test_asyncresult.py

@@ -1,266 +1,267 b''
 """Tests for asyncresult.py
 
 Authors:
 
 * Min RK
 """
 
 #-------------------------------------------------------------------------------
 #  Copyright (C) 2011  The IPython Development Team
 #
 #  Distributed under the terms of the BSD License.  The full license is in
 #  the file COPYING, distributed as part of this software.
 #-------------------------------------------------------------------------------
 
 #-------------------------------------------------------------------------------
 # Imports
 #-------------------------------------------------------------------------------
 
 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)
 
 def wait(n):
     import time
     time.sleep(n)
     return n
 
 class AsyncResultTest(ClusterTestCase):
 
     def test_single_result_view(self):
         """various one-target views get the right value for single_result"""
         eid = self.client.ids[-1]
         ar = self.client[eid].apply_async(lambda : 42)
         self.assertEquals(ar.get(), 42)
         ar = self.client[[eid]].apply_async(lambda : 42)
         self.assertEquals(ar.get(), [42])
         ar = self.client[-1:].apply_async(lambda : 42)
         self.assertEquals(ar.get(), [42])
 
     def test_get_after_done(self):
         ar = self.client[-1].apply_async(lambda : 42)
         ar.wait()
         self.assertTrue(ar.ready())
         self.assertEquals(ar.get(), 42)
         self.assertEquals(ar.get(), 42)
 
     def test_get_before_done(self):
         ar = self.client[-1].apply_async(wait, 0.1)
         self.assertRaises(TimeoutError, ar.get, 0)
         ar.wait(0)
         self.assertFalse(ar.ready())
         self.assertEquals(ar.get(), 0.1)
 
     def test_get_after_error(self):
         ar = self.client[-1].apply_async(lambda : 1/0)
         ar.wait(10)
         self.assertRaisesRemote(ZeroDivisionError, ar.get)
         self.assertRaisesRemote(ZeroDivisionError, ar.get)
         self.assertRaisesRemote(ZeroDivisionError, ar.get_dict)
     
     def test_get_dict(self):
         n = len(self.client)
         ar = self.client[:].apply_async(lambda : 5)
         self.assertEquals(ar.get(), [5]*n)
         d = ar.get_dict()
         self.assertEquals(sorted(d.keys()), sorted(self.client.ids))
         for eid,r in d.iteritems():
             self.assertEquals(r, 5)
     
     def test_list_amr(self):
         ar = self.client.load_balanced_view().map_async(wait, [0.1]*5)
         rlist = list(ar)
     
     def test_getattr(self):
         ar = self.client[:].apply_async(wait, 0.5)
         self.assertRaises(AttributeError, lambda : ar._foo)
         self.assertRaises(AttributeError, lambda : ar.__length_hint__())
         self.assertRaises(AttributeError, lambda : ar.foo)
         self.assertRaises(AttributeError, lambda : ar.engine_id)
         self.assertFalse(hasattr(ar, '__length_hint__'))
         self.assertFalse(hasattr(ar, 'foo'))
         self.assertFalse(hasattr(ar, 'engine_id'))
         ar.get(5)
         self.assertRaises(AttributeError, lambda : ar._foo)
         self.assertRaises(AttributeError, lambda : ar.__length_hint__())
         self.assertRaises(AttributeError, lambda : ar.foo)
         self.assertTrue(isinstance(ar.engine_id, list))
         self.assertEquals(ar.engine_id, ar['engine_id'])
         self.assertFalse(hasattr(ar, '__length_hint__'))
         self.assertFalse(hasattr(ar, 'foo'))
         self.assertTrue(hasattr(ar, 'engine_id'))
 
     def test_getitem(self):
         ar = self.client[:].apply_async(wait, 0.5)
         self.assertRaises(TimeoutError, lambda : ar['foo'])
         self.assertRaises(TimeoutError, lambda : ar['engine_id'])
         ar.get(5)
         self.assertRaises(KeyError, lambda : ar['foo'])
         self.assertTrue(isinstance(ar['engine_id'], list))
         self.assertEquals(ar.engine_id, ar['engine_id'])
     
     def test_single_result(self):
         ar = self.client[-1].apply_async(wait, 0.5)
         self.assertRaises(TimeoutError, lambda : ar['foo'])
         self.assertRaises(TimeoutError, lambda : ar['engine_id'])
         self.assertTrue(ar.get(5) == 0.5)
         self.assertTrue(isinstance(ar['engine_id'], int))
         self.assertTrue(isinstance(ar.engine_id, int))
         self.assertEquals(ar.engine_id, ar['engine_id'])
     
     def test_abort(self):
         e = self.client[-1]
         ar = e.execute('import time; time.sleep(1)', block=False)
         ar2 = e.apply_async(lambda : 2)
         ar2.abort()
         self.assertRaises(error.TaskAborted, ar2.get)
         ar.get()
     
     def test_len(self):
         v = self.client.load_balanced_view()
         ar = v.map_async(lambda x: x, range(10))
         self.assertEquals(len(ar), 10)
         ar = v.apply_async(lambda x: x, range(10))
         self.assertEquals(len(ar), 1)
         ar = self.client[:].apply_async(lambda x: x, range(10))
         self.assertEquals(len(ar), len(self.client.ids))
     
     def test_wall_time_single(self):
         v = self.client.load_balanced_view()
         ar = v.apply_async(time.sleep, 0.25)
         self.assertRaises(TimeoutError, getattr, ar, 'wall_time')
         ar.get(2)
         self.assertTrue(ar.wall_time < 1.)
         self.assertTrue(ar.wall_time > 0.2)
 
     def test_wall_time_multi(self):
         self.minimum_engines(4)
         v = self.client[:]
         ar = v.apply_async(time.sleep, 0.25)
         self.assertRaises(TimeoutError, getattr, ar, 'wall_time')
         ar.get(2)
         self.assertTrue(ar.wall_time < 1.)
         self.assertTrue(ar.wall_time > 0.2)
 
     def test_serial_time_single(self):
         v = self.client.load_balanced_view()
         ar = v.apply_async(time.sleep, 0.25)
         self.assertRaises(TimeoutError, getattr, ar, 'serial_time')
         ar.get(2)
         self.assertTrue(ar.serial_time < 1.)
         self.assertTrue(ar.serial_time > 0.2)
 
     def test_serial_time_multi(self):
         self.minimum_engines(4)
         v = self.client[:]
         ar = v.apply_async(time.sleep, 0.25)
         self.assertRaises(TimeoutError, getattr, ar, 'serial_time')
         ar.get(2)
         self.assertTrue(ar.serial_time < 2.)
         self.assertTrue(ar.serial_time > 0.8)
 
     def test_elapsed_single(self):
         v = self.client.load_balanced_view()
         ar = v.apply_async(time.sleep, 0.25)
         while not ar.ready():
             time.sleep(0.01)
             self.assertTrue(ar.elapsed < 1)
         self.assertTrue(ar.elapsed < 1)
         ar.get(2)
 
     def test_elapsed_multi(self):
         v = self.client[:]
         ar = v.apply_async(time.sleep, 0.25)
         while not ar.ready():
             time.sleep(0.01)
             self.assertTrue(ar.elapsed < 1)
         self.assertTrue(ar.elapsed < 1)
         ar.get(2)
 
     def test_hubresult_timestamps(self):
         self.minimum_engines(4)
         v = self.client[:]
         ar = v.apply_async(time.sleep, 0.25)
         ar.get(2)
         rc2 = Client(profile='iptest')
         # must have try/finally to close second Client, otherwise
         # will have dangling sockets causing problems
         try:
             time.sleep(0.25)
             hr = rc2.get_result(ar.msg_ids)
             self.assertTrue(hr.elapsed > 0., "got bad elapsed: %s" % hr.elapsed)
             hr.get(1)
             self.assertTrue(hr.wall_time < ar.wall_time + 0.2, "got bad wall_time: %s > %s" % (hr.wall_time, ar.wall_time))
             self.assertEquals(hr.serial_time, ar.serial_time)
         finally:
             rc2.close()
 
     def test_display_empty_streams_single(self):
         """empty stdout/err are not displayed (single result)"""
         self.minimum_engines(1)
         
         v = self.client[-1]
         ar = v.execute("print (5555)")
         ar.get(5)
         with capture_output() as io:
             ar.display_outputs()
         self.assertEquals(io.stderr, '')
         self.assertEquals('5555\n', io.stdout)
 
         ar = v.execute("a=5")
         ar.get(5)
         with capture_output() as io:
             ar.display_outputs()
         self.assertEquals(io.stderr, '')
         self.assertEquals(io.stdout, '')
         
     def test_display_empty_streams_type(self):
         """empty stdout/err are not displayed (groupby type)"""
         self.minimum_engines(1)
         
         v = self.client[:]
         ar = v.execute("print (5555)")
         ar.get(5)
         with capture_output() as io:
             ar.display_outputs()
         self.assertEquals(io.stderr, '')
         self.assertEquals(io.stdout.count('5555'), len(v), io.stdout)
         self.assertFalse('\n\n' in io.stdout, io.stdout)
         self.assertEquals(io.stdout.count('[stdout:'), len(v), io.stdout)
 
         ar = v.execute("a=5")
         ar.get(5)
         with capture_output() as io:
             ar.display_outputs()
         self.assertEquals(io.stderr, '')
         self.assertEquals(io.stdout, '')
 
     def test_display_empty_streams_engine(self):
         """empty stdout/err are not displayed (groupby engine)"""
         self.minimum_engines(1)
         
         v = self.client[:]
         ar = v.execute("print (5555)")
         ar.get(5)
         with capture_output() as io:
             ar.display_outputs('engine')
         self.assertEquals(io.stderr, '')
         self.assertEquals(io.stdout.count('5555'), len(v), io.stdout)
         self.assertFalse('\n\n' in io.stdout, io.stdout)
         self.assertEquals(io.stdout.count('[stdout:'), len(v), io.stdout)
 
         ar = v.execute("a=5")
         ar.get(5)
         with capture_output() as io:
             ar.display_outputs('engine')
         self.assertEquals(io.stderr, '')
         self.assertEquals(io.stdout, '')
         
 

IPython/parallel/tests/test_magics.py

@@ -1,339 +1,340 b''
 # -*- coding: utf-8 -*-
 """Test Parallel magics
 
 Authors:
 
 * Min RK
 """
 #-------------------------------------------------------------------------------
 #  Copyright (C) 2011  The IPython Development Team
 #
 #  Distributed under the terms of the BSD License.  The full license is in
 #  the file COPYING, distributed as part of this software.
 #-------------------------------------------------------------------------------
 
 #-------------------------------------------------------------------------------
 # Imports
 #-------------------------------------------------------------------------------
 
 import re
 import sys
 import time
 
 import zmq
 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
 from IPython.parallel import AsyncResult
 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)
 
 class TestParallelMagics(ClusterTestCase, ParametricTestCase):
     
     def test_px_blocking(self):
         ip = get_ipython()
         v = self.client[-1:]
         v.activate()
         v.block=True
 
         ip.magic('px a=5')
         self.assertEquals(v['a'], [5])
         ip.magic('px a=10')
         self.assertEquals(v['a'], [10])
         # just 'print a' works ~99% of the time, but this ensures that
         # the stdout message has arrived when the result is finished:
         with capture_output() as io:
             ip.magic(
                 'px import sys,time;print(a);sys.stdout.flush();time.sleep(0.2)'
             )
         out = io.stdout
         self.assertTrue('[stdout:' in out, out)
         self.assertFalse('\n\n' in out)
         self.assertTrue(out.rstrip().endswith('10'))
         self.assertRaisesRemote(ZeroDivisionError, ip.magic, 'px 1/0')
     
     def _check_generated_stderr(self, stderr, n):
         expected = [
             r'\[stderr:\d+\]',
             '^stderr$',
             '^stderr2$',
         ] * n
         
         self.assertFalse('\n\n' in stderr, stderr)
         lines = stderr.splitlines()
         self.assertEquals(len(lines), len(expected), stderr)
         for line,expect in zip(lines, expected):
             if isinstance(expect, str):
                 expect = [expect]
             for ex in expect:
                 self.assertTrue(re.search(ex, line) is not None, "Expected %r in %r" % (ex, line))
         
     def test_cellpx_block_args(self):
         """%%px --[no]block flags work"""
         ip = get_ipython()
         v = self.client[-1:]
         v.activate()
         v.block=False
         
         for block in (True, False):
             v.block = block
             
             with capture_output() as io:
                 ip.run_cell_magic("px", "", "1")
             if block:
                 self.assertTrue(io.stdout.startswith("Parallel"), io.stdout)
             else:
                 self.assertTrue(io.stdout.startswith("Async"), io.stdout)
             
             with capture_output() as io:
                 ip.run_cell_magic("px", "--block", "1")
             self.assertTrue(io.stdout.startswith("Parallel"), io.stdout)
 
             with capture_output() as io:
                 ip.run_cell_magic("px", "--noblock", "1")
             self.assertTrue(io.stdout.startswith("Async"), io.stdout)
     
     def test_cellpx_groupby_engine(self):
         """%%px --group-outputs=engine"""
         ip = get_ipython()
         v = self.client[:]
         v.block = True
         v.activate()
         
         v['generate_output'] = generate_output
         
         with capture_output() as io:
             ip.run_cell_magic('px', '--group-outputs=engine', 'generate_output()')
         
         self.assertFalse('\n\n' in io.stdout)
         lines = io.stdout.splitlines()[1:]
         expected = [
             r'\[stdout:\d+\]',
             'stdout',
             'stdout2',
             r'\[output:\d+\]',
             r'IPython\.core\.display\.HTML',
             r'IPython\.core\.display\.Math',
             r'Out\[\d+:\d+\]:.*IPython\.core\.display\.Math',
             ] * len(v)
         
         self.assertEquals(len(lines), len(expected), io.stdout)
         for line,expect in zip(lines, expected):
             if isinstance(expect, str):
                 expect = [expect]
             for ex in expect:
                 self.assertTrue(re.search(ex, line) is not None, "Expected %r in %r" % (ex, line))
         
         self._check_generated_stderr(io.stderr, len(v))
 
 
     def test_cellpx_groupby_order(self):
         """%%px --group-outputs=order"""
         ip = get_ipython()
         v = self.client[:]
         v.block = True
         v.activate()
         
         v['generate_output'] = generate_output
         
         with capture_output() as io:
             ip.run_cell_magic('px', '--group-outputs=order', 'generate_output()')
         
         self.assertFalse('\n\n' in io.stdout)
         lines = io.stdout.splitlines()[1:]
         expected = []
         expected.extend([
             r'\[stdout:\d+\]',
             'stdout',
             'stdout2',
         ] * len(v))
         expected.extend([
             r'\[output:\d+\]',
             'IPython.core.display.HTML',
         ] * len(v))
         expected.extend([
             r'\[output:\d+\]',
             'IPython.core.display.Math',
         ] * len(v))
         expected.extend([
             r'Out\[\d+:\d+\]:.*IPython\.core\.display\.Math'
         ] * len(v))
         
         self.assertEquals(len(lines), len(expected), io.stdout)
         for line,expect in zip(lines, expected):
             if isinstance(expect, str):
                 expect = [expect]
             for ex in expect:
                 self.assertTrue(re.search(ex, line) is not None, "Expected %r in %r" % (ex, line))
         
         self._check_generated_stderr(io.stderr, len(v))
 
     def test_cellpx_groupby_type(self):
         """%%px --group-outputs=type"""
         ip = get_ipython()
         v = self.client[:]
         v.block = True
         v.activate()
         
         v['generate_output'] = generate_output
         
         with capture_output() as io:
             ip.run_cell_magic('px', '--group-outputs=type', 'generate_output()')
         
         self.assertFalse('\n\n' in io.stdout)
         lines = io.stdout.splitlines()[1:]
         
         expected = []
         expected.extend([
             r'\[stdout:\d+\]',
             'stdout',
             'stdout2',
         ] * len(v))
         expected.extend([
             r'\[output:\d+\]',
             r'IPython\.core\.display\.HTML',
             r'IPython\.core\.display\.Math',
         ] * len(v))
         expected.extend([
             (r'Out\[\d+:\d+\]', r'IPython\.core\.display\.Math')
         ] * len(v))
         
         self.assertEquals(len(lines), len(expected), io.stdout)
         for line,expect in zip(lines, expected):
             if isinstance(expect, str):
                 expect = [expect]
             for ex in expect:
                 self.assertTrue(re.search(ex, line) is not None, "Expected %r in %r" % (ex, line))
         
         self._check_generated_stderr(io.stderr, len(v))
 
 
     def test_px_nonblocking(self):
         ip = get_ipython()
         v = self.client[-1:]
         v.activate()
         v.block=False
 
         ip.magic('px a=5')
         self.assertEquals(v['a'], [5])
         ip.magic('px a=10')
         self.assertEquals(v['a'], [10])
         with capture_output() as io:
             ar = ip.magic('px print (a)')
         self.assertTrue(isinstance(ar, AsyncResult))
         self.assertTrue('Async' in io.stdout)
         self.assertFalse('[stdout:' in io.stdout)
         self.assertFalse('\n\n' in io.stdout)
         
         ar = ip.magic('px 1/0')
         self.assertRaisesRemote(ZeroDivisionError, ar.get)
     
     def test_autopx_blocking(self):
         ip = get_ipython()
         v = self.client[-1]
         v.activate()
         v.block=True
 
         with capture_output() as io:
             ip.magic('autopx')
             ip.run_cell('\n'.join(('a=5','b=12345','c=0')))
             ip.run_cell('b*=2')
             ip.run_cell('print (b)')
             ip.run_cell('b')
             ip.run_cell("b/c")
             ip.magic('autopx')
         
         output = io.stdout
         
         self.assertTrue(output.startswith('%autopx enabled'), output)
         self.assertTrue(output.rstrip().endswith('%autopx disabled'), output)
         self.assertTrue('RemoteError: ZeroDivisionError' in output, output)
         self.assertTrue('\nOut[' in output, output)
         self.assertTrue(': 24690' in output, output)
         ar = v.get_result(-1)
         self.assertEquals(v['a'], 5)
         self.assertEquals(v['b'], 24690)
         self.assertRaisesRemote(ZeroDivisionError, ar.get)
 
     def test_autopx_nonblocking(self):
         ip = get_ipython()
         v = self.client[-1]
         v.activate()
         v.block=False
 
         with capture_output() as io:
             ip.magic('autopx')
             ip.run_cell('\n'.join(('a=5','b=10','c=0')))
             ip.run_cell('print (b)')
             ip.run_cell('import time; time.sleep(0.1)')
             ip.run_cell("b/c")
             ip.run_cell('b*=2')
             ip.magic('autopx')
         
         output = io.stdout.rstrip()
         
         self.assertTrue(output.startswith('%autopx enabled'))
         self.assertTrue(output.endswith('%autopx disabled'))
         self.assertFalse('ZeroDivisionError' in output)
         ar = v.get_result(-2)
         self.assertRaisesRemote(ZeroDivisionError, ar.get)
         # prevent TaskAborted on pulls, due to ZeroDivisionError
         time.sleep(0.5)
         self.assertEquals(v['a'], 5)
         # b*=2 will not fire, due to abort
         self.assertEquals(v['b'], 10)
     
     def test_result(self):
         ip = get_ipython()
         v = self.client[-1]
         v.activate()
         data = dict(a=111,b=222)
         v.push(data, block=True)
 
         ip.magic('px a')
         ip.magic('px b')
         for idx, name in [
                     ('', 'b'),
                     ('-1', 'b'),
                     ('2', 'b'),
                     ('1', 'a'),
                     ('-2', 'a'),
                 ]:
             with capture_output() as io:
                 ip.magic('result ' + idx)
             output = io.stdout
             msg = "expected %s output to include %s, but got: %s" % \
                 ('%result '+idx, str(data[name]), output)
             self.assertTrue(str(data[name]) in output, msg)
         
     @dec.skipif_not_matplotlib
     def test_px_pylab(self):
         """%pylab works on engines"""
         ip = get_ipython()
         v = self.client[-1]
         v.block = True
         v.activate()
         
         with capture_output() as io:
             ip.magic("px %pylab inline")
         
         self.assertTrue("Welcome to pylab" in io.stdout, io.stdout)
         self.assertTrue("backend_inline" in io.stdout, io.stdout)
         
         with capture_output() as io:
             ip.magic("px plot(rand(100))")
         
         self.assertTrue('Out[' in io.stdout, io.stdout)
         self.assertTrue('matplotlib.lines' in io.stdout, io.stdout)
         
 

IPython/utils/io.py

@@ -1,323 +1,357 b''
 # encoding: utf-8
 """
 IO related utilities.
 """
 
 #-----------------------------------------------------------------------------
 #  Copyright (C) 2008-2011  The IPython Development Team
 #
 #  Distributed under the terms of the BSD License.  The full license is in
 #  the file COPYING, distributed as part of this software.
 #-----------------------------------------------------------------------------
 from __future__ import print_function
 
 #-----------------------------------------------------------------------------
 # Imports
 #-----------------------------------------------------------------------------
 import os
 import sys
 import tempfile
+from StringIO import StringIO
 
 #-----------------------------------------------------------------------------
 # Code
 #-----------------------------------------------------------------------------
 
 
 class IOStream:
 
     def __init__(self,stream, fallback=None):
         if not hasattr(stream,'write') or not hasattr(stream,'flush'):
             if fallback is not None:
                 stream = fallback
             else:
                 raise ValueError("fallback required, but not specified")
         self.stream = stream
         self._swrite = stream.write
 
         # clone all methods not overridden:
         def clone(meth):
             return not hasattr(self, meth) and not meth.startswith('_')
         for meth in filter(clone, dir(stream)):
             setattr(self, meth, getattr(stream, meth))
 
     def write(self,data):
         try:
             self._swrite(data)
         except:
             try:
                 # print handles some unicode issues which may trip a plain
                 # write() call.  Emulate write() by using an empty end
                 # argument.
                 print(data, end='', file=self.stream)
             except:
                 # if we get here, something is seriously broken.
                 print('ERROR - failed to write data to stream:', self.stream,
                       file=sys.stderr)
 
     def writelines(self, lines):
         if isinstance(lines, basestring):
             lines = [lines]
         for line in lines:
             self.write(line)
 
     # This class used to have a writeln method, but regular files and streams
     # in Python don't have this method. We need to keep this completely
     # compatible so we removed it.
 
     @property
     def closed(self):
         return self.stream.closed
 
     def close(self):
         pass
 
 # setup stdin/stdout/stderr to sys.stdin/sys.stdout/sys.stderr
 devnull = open(os.devnull, 'a')
 stdin = IOStream(sys.stdin, fallback=devnull)
 stdout = IOStream(sys.stdout, fallback=devnull)
 stderr = IOStream(sys.stderr, fallback=devnull)
 
 class IOTerm:
     """ Term holds the file or file-like objects for handling I/O operations.
 
     These are normally just sys.stdin, sys.stdout and sys.stderr but for
     Windows they can can replaced to allow editing the strings before they are
     displayed."""
 
     # In the future, having IPython channel all its I/O operations through
     # this class will make it easier to embed it into other environments which
     # are not a normal terminal (such as a GUI-based shell)
     def __init__(self, stdin=None, stdout=None, stderr=None):
         mymodule = sys.modules[__name__]
         self.stdin  = IOStream(stdin, mymodule.stdin)
         self.stdout = IOStream(stdout, mymodule.stdout)
         self.stderr = IOStream(stderr, mymodule.stderr)
 
 
 class Tee(object):
     """A class to duplicate an output stream to stdout/err.
 
     This works in a manner very similar to the Unix 'tee' command.
 
     When the object is closed or deleted, it closes the original file given to
     it for duplication.
     """
     # Inspired by:
     # http://mail.python.org/pipermail/python-list/2007-May/442737.html
 
     def __init__(self, file_or_name, mode="w", channel='stdout'):
         """Construct a new Tee object.
 
         Parameters
         ----------
         file_or_name : filename or open filehandle (writable)
           File that will be duplicated
 
         mode : optional, valid mode for open().
           If a filename was give, open with this mode.
 
         channel : str, one of ['stdout', 'stderr']
         """
         if channel not in ['stdout', 'stderr']:
             raise ValueError('Invalid channel spec %s' % channel)
 
         if hasattr(file_or_name, 'write') and hasattr(file_or_name, 'seek'):
             self.file = file_or_name
         else:
             self.file = open(file_or_name, mode)
         self.channel = channel
         self.ostream = getattr(sys, channel)
         setattr(sys, channel, self)
         self._closed = False
 
     def close(self):
         """Close the file and restore the channel."""
         self.flush()
         setattr(sys, self.channel, self.ostream)
         self.file.close()
         self._closed = True
 
     def write(self, data):
         """Write data to both channels."""
         self.file.write(data)
         self.ostream.write(data)
         self.ostream.flush()
 
     def flush(self):
         """Flush both channels."""
         self.file.flush()
         self.ostream.flush()
 
     def __del__(self):
         if not self._closed:
             self.close()
 
 
 def file_read(filename):
     """Read a file and close it.  Returns the file source."""
     fobj = open(filename,'r');
     source = fobj.read();
     fobj.close()
     return source
 
 
 def file_readlines(filename):
     """Read a file and close it.  Returns the file source using readlines()."""
     fobj = open(filename,'r');
     lines = fobj.readlines();
     fobj.close()
     return lines
 
 
 def raw_input_multi(header='', ps1='==> ', ps2='..> ',terminate_str = '.'):
     """Take multiple lines of input.
 
     A list with each line of input as a separate element is returned when a
     termination string is entered (defaults to a single '.'). Input can also
     terminate via EOF (^D in Unix, ^Z-RET in Windows).
 
     Lines of input which end in \\ are joined into single entries (and a
     secondary continuation prompt is issued as long as the user terminates
     lines with \\). This allows entering very long strings which are still
     meant to be treated as single entities.
     """
 
     try:
         if header:
             header += '\n'
         lines = [raw_input(header + ps1)]
     except EOFError:
         return []
     terminate = [terminate_str]
     try:
         while lines[-1:] != terminate:
             new_line = raw_input(ps1)
             while new_line.endswith('\\'):
                 new_line = new_line[:-1] + raw_input(ps2)
             lines.append(new_line)
 
         return lines[:-1]  # don't return the termination command
     except EOFError:
         print()
         return lines
 
 
 def raw_input_ext(prompt='',  ps2='... '):
     """Similar to raw_input(), but accepts extended lines if input ends with \\."""
 
     line = raw_input(prompt)
     while line.endswith('\\'):
         line = line[:-1] + raw_input(ps2)
     return line
 
 
 def ask_yes_no(prompt,default=None):
     """Asks a question and returns a boolean (y/n) answer.
 
     If default is given (one of 'y','n'), it is used if the user input is
     empty. Otherwise the question is repeated until an answer is given.
 
     An EOF is treated as the default answer.  If there is no default, an
     exception is raised to prevent infinite loops.
 
     Valid answers are: y/yes/n/no (match is not case sensitive)."""
 
     answers = {'y':True,'n':False,'yes':True,'no':False}
     ans = None
     while ans not in answers.keys():
         try:
             ans = raw_input(prompt+' ').lower()
             if not ans:  # response was an empty string
                 ans = default
         except KeyboardInterrupt:
             pass
         except EOFError:
             if default in answers.keys():
                 ans = default
                 print()
             else:
                 raise
 
     return answers[ans]
 
 
 class NLprinter:
     """Print an arbitrarily nested list, indicating index numbers.
 
     An instance of this class called nlprint is available and callable as a
     function.
 
     nlprint(list,indent=' ',sep=': ') -> prints indenting each level by 'indent'
     and using 'sep' to separate the index from the value. """
 
     def __init__(self):
         self.depth = 0
 
     def __call__(self,lst,pos='',**kw):
         """Prints the nested list numbering levels."""
         kw.setdefault('indent',' ')
         kw.setdefault('sep',': ')
         kw.setdefault('start',0)
         kw.setdefault('stop',len(lst))
         # we need to remove start and stop from kw so they don't propagate
         # into a recursive call for a nested list.
         start = kw['start']; del kw['start']
         stop = kw['stop']; del kw['stop']
         if self.depth == 0 and 'header' in kw.keys():
             print(kw['header'])
 
         for idx in range(start,stop):
             elem = lst[idx]
             newpos = pos + str(idx)
             if type(elem)==type([]):
                 self.depth += 1
                 self.__call__(elem, newpos+",", **kw)
                 self.depth -= 1
             else:
                 print(kw['indent']*self.depth + newpos + kw["sep"] + repr(elem))
 
 nlprint = NLprinter()
 
 
 def temp_pyfile(src, ext='.py'):
     """Make a temporary python file, return filename and filehandle.
 
     Parameters
     ----------
     src : string or list of strings (no need for ending newlines if list)
       Source code to be written to the file.
 
     ext : optional, string
       Extension for the generated file.
 
     Returns
     -------
     (filename, open filehandle)
       It is the caller's responsibility to close the open file and unlink it.
     """
     fname = tempfile.mkstemp(ext)[1]
     f = open(fname,'w')
     f.write(src)
     f.flush()
     return fname, f
 
 
 def raw_print(*args, **kw):
     """Raw print to sys.__stdout__, otherwise identical interface to print()."""
 
     print(*args, sep=kw.get('sep', ' '), end=kw.get('end', '\n'),
           file=sys.__stdout__)
     sys.__stdout__.flush()
 
 
 def raw_print_err(*args, **kw):
     """Raw print to sys.__stderr__, otherwise identical interface to print()."""
 
     print(*args, sep=kw.get('sep', ' '), end=kw.get('end', '\n'),
           file=sys.__stderr__)
     sys.__stderr__.flush()
 
 
 # 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
+
+

IPython/utils/tests/test_io.py

@@ -1,75 +1,85 b''
 # encoding: utf-8
 """Tests for io.py"""
 
 #-----------------------------------------------------------------------------
 #  Copyright (C) 2008-2011  The IPython Development Team
 #
 #  Distributed under the terms of the BSD License.  The full license is in
 #  the file COPYING, distributed as part of this software.
 #-----------------------------------------------------------------------------
 
 #-----------------------------------------------------------------------------
 # Imports
 #-----------------------------------------------------------------------------
 
 import sys
 
 from StringIO import StringIO
 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
 
 #-----------------------------------------------------------------------------
 # Tests
 #-----------------------------------------------------------------------------
 
 
 def test_tee_simple():
     "Very simple check with stdout only"
     chan = StringIO()
     text = 'Hello'
     tee = Tee(chan, channel='stdout')
     print >> chan, text
     nt.assert_equal(chan.getvalue(), text+"\n")
 
 
 class TeeTestCase(dec.ParametricTestCase):
 
     def tchan(self, channel, check='close'):
         trap = StringIO()
         chan = StringIO()
         text = 'Hello'
         
         std_ori = getattr(sys, channel)
         setattr(sys, channel, trap)
 
         tee = Tee(chan, channel=channel)
         print >> chan, text,
         setattr(sys, channel, std_ori)
         trap_val = trap.getvalue()
         nt.assert_equals(chan.getvalue(), text)
         if check=='close':
             tee.close()
         else:
             del tee
 
     def test(self):
         for chan in ['stdout', 'stderr']:
             for check in ['close', 'del']:
                 yield self.tchan(chan, check)
 
 def test_io_init():
     """Test that io.stdin/out/err exist at startup"""
     for name in ('stdin', 'stdout', 'stderr'):
         cmd = doctest_refactor_print("from IPython.utils import io;print io.%s.__class__"%name)
         p = Popen([sys.executable, '-c', cmd],
                     stdout=PIPE)
         p.wait()
         classname = p.stdout.read().strip().decode('ascii')
         # __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')