Merge pull request #1870 from minrk/captureio

New `%%capture` cell magic captures stdout/err while running a cell.

Uses `capture_output()` context manager, moved to utils.io from IPython.parallel testing utilities, where it originated.

The caputre objects can be printed as a string, case in which they display the captured stdout, which is also available as `.stdout`.  The captured stderr, if any, is in a `.stderr` attribute.  A `.show()` method can be called to quickly print both, with stderr being correctly printed to the sys.stderr stream (so the notebook displays it with red highlighting).

closes  #1863

docs/examples/notebooks/Capturing Output.ipynb

@@ -0,0 +1,182 b''
+{
+ "metadata": {
+  "name": "Capturing Output"
+ },
+ "nbformat": 3,
+ "worksheets": [
+  {
+   "cells": [
+    {
+     "cell_type": "heading",
+     "level": 1,
+     "source": [
+      "Capturing Output with <tt>%%capture</tt>"
+     ]
+    },
+    {
+     "cell_type": "markdown",
+     "source": [
+      "One of IPython's new cell magics is `%%capture`, which captures stdout/err for a cell,",
+      "and discards them or stores them in variables in your namespace."
+     ]
+    },
+    {
+     "cell_type": "code",
+     "input": [
+      "import sys"
+     ],
+     "language": "python",
+     "outputs": []
+    },
+    {
+     "cell_type": "markdown",
+     "source": [
+      "By default, it just swallows it up.  This is a simple way to suppress unwanted output."
+     ]
+    },
+    {
+     "cell_type": "code",
+     "input": [
+      "%%capture",
+      "print 'hi, stdout'",
+      "print >> sys.stderr, 'hi, stderr'"
+     ],
+     "language": "python",
+     "outputs": []
+    },
+    {
+     "cell_type": "markdown",
+     "source": [
+      "If you specify a name, then stdout and stderr will be stored in an object in your namespace."
+     ]
+    },
+    {
+     "cell_type": "code",
+     "input": [
+      "%%capture captured",
+      "print 'hi, stdout'",
+      "print >> sys.stderr, 'hi, stderr'"
+     ],
+     "language": "python",
+     "outputs": []
+    },
+    {
+     "cell_type": "code",
+     "input": [
+      "captured"
+     ],
+     "language": "python",
+     "outputs": []
+    },
+    {
+     "cell_type": "markdown",
+     "source": [
+      "Calling the object writes the output to stdout/err as appropriate."
+     ]
+    },
+    {
+     "cell_type": "code",
+     "input": [
+      "captured()"
+     ],
+     "language": "python",
+     "outputs": []
+    },
+    {
+     "cell_type": "code",
+     "input": [
+      "captured.stdout"
+     ],
+     "language": "python",
+     "outputs": []
+    },
+    {
+     "cell_type": "code",
+     "input": [
+      "captured.stderr"
+     ],
+     "language": "python",
+     "outputs": []
+    },
+    {
+     "cell_type": "markdown",
+     "source": [
+      "`%%capture` only captures stdout/err, not displaypub, so you can still do plots and use the display protocol inside %%capture"
+     ]
+    },
+    {
+     "cell_type": "code",
+     "input": [
+      "%pylab inline"
+     ],
+     "language": "python",
+     "outputs": []
+    },
+    {
+     "cell_type": "code",
+     "input": [
+      "%%capture wontshutup",
+      "",
+      "print \"setting up X\"",
+      "x = np.linspace(0,5,1000)",
+      "print \"step 2: constructing y-data\"",
+      "y = np.sin(x)",
+      "print \"step 3: display info about y\"",
+      "plt.plot(x,y)",
+      "print \"okay, I'm done now\""
+     ],
+     "language": "python",
+     "outputs": []
+    },
+    {
+     "cell_type": "code",
+     "input": [
+      "wontshutup()"
+     ],
+     "language": "python",
+     "outputs": []
+    },
+    {
+     "cell_type": "markdown",
+     "source": [
+      "And you can selectively disable capturing stdout or stderr by passing `--no-stdout/err`."
+     ]
+    },
+    {
+     "cell_type": "code",
+     "input": [
+      "%%capture cap --no-stderr",
+      "print 'hi, stdout'",
+      "print >> sys.stderr, \"hello, stderr\""
+     ],
+     "language": "python",
+     "outputs": []
+    },
+    {
+     "cell_type": "code",
+     "input": [
+      "cap.stdout"
+     ],
+     "language": "python",
+     "outputs": []
+    },
+    {
+     "cell_type": "code",
+     "input": [
+      "cap.stderr"
+     ],
+     "language": "python",
+     "outputs": []
+    },
+    {
+     "cell_type": "code",
+     "input": [
+      ""
+     ],
+     "language": "python",
+     "outputs": []
+    }
+   ]
+  }
+ ]
+}
  No newline at end of file

IPython/core/magics/execution.py

@@ -1,990 +1,1022 b''
 """Implementation of execution-related magic functions.
 """
 #-----------------------------------------------------------------------------
 #  Copyright (c) 2012 The IPython Development Team.
 #
 #  Distributed under the terms of the Modified BSD License.
 #
 #  The full license is in the file COPYING.txt, distributed with this software.
 #-----------------------------------------------------------------------------
 
 #-----------------------------------------------------------------------------
 # Imports
 #-----------------------------------------------------------------------------
 
 # Stdlib
 import __builtin__ as builtin_mod
 import bdb
 import os
 import sys
 import time
 from StringIO import StringIO
 
 # cProfile was added in Python2.5
 try:
     import cProfile as profile
     import pstats
 except ImportError:
     # profile isn't bundled by default in Debian for license reasons
     try:
         import profile, pstats
     except ImportError:
         profile = pstats = None
 
 # Our own packages
 from IPython.core import debugger, oinspect
+from IPython.core import magic_arguments
 from IPython.core import page
 from IPython.core.error import UsageError
 from IPython.core.macro import Macro
-from IPython.core.magic import (Magics, magics_class, line_magic, 
+from IPython.core.magic import (Magics, magics_class, line_magic, cell_magic,
                                 line_cell_magic, on_off, needs_local_scope)
 from IPython.testing.skipdoctest import skip_doctest
 from IPython.utils import py3compat
+from IPython.utils.io import capture_output
 from IPython.utils.ipstruct import Struct
 from IPython.utils.module_paths import find_mod
 from IPython.utils.path import get_py_filename, unquote_filename
 from IPython.utils.timing import clock, clock2
 from IPython.utils.warn import warn, error
 
 #-----------------------------------------------------------------------------
 # Magic implementation classes
 #-----------------------------------------------------------------------------
 
 @magics_class
 class ExecutionMagics(Magics):
     """Magics related to code execution, debugging, profiling, etc.
 
     """
 
     def __init__(self, shell):
         super(ExecutionMagics, self).__init__(shell)
         if profile is None:
             self.prun = self.profile_missing_notice
         # Default execution function used to actually run user code.
         self.default_runner = None
 
     def profile_missing_notice(self, *args, **kwargs):
         error("""\
 The profile module could not be found. It has been removed from the standard
 python packages because of its non-free license. To use profiling, install the
 python-profiler package from non-free.""")
 
     @skip_doctest
     @line_cell_magic
     def prun(self, parameter_s='', cell=None, user_mode=True,
                    opts=None,arg_lst=None,prog_ns=None):
 
         """Run a statement through the python code profiler.
 
         Usage, in line mode:
           %prun [options] statement
 
         Usage, in cell mode:
           %%prun [options] [statement]
           code...
           code...
 
         In cell mode, the additional code lines are appended to the (possibly
         empty) statement in the first line.  Cell mode allows you to easily
         profile multiline blocks without having to put them in a separate
         function.
         
         The given statement (which doesn't require quote marks) is run via the
         python profiler in a manner similar to the profile.run() function.
         Namespaces are internally managed to work correctly; profile.run
         cannot be used in IPython because it makes certain assumptions about
         namespaces which do not hold under IPython.
 
         Options:
 
         -l <limit>: you can place restrictions on what or how much of the
         profile gets printed. The limit value can be:
 
           * A string: only information for function names containing this string
           is printed.
 
           * An integer: only these many lines are printed.
 
           * A float (between 0 and 1): this fraction of the report is printed
           (for example, use a limit of 0.4 to see the topmost 40% only).
 
         You can combine several limits with repeated use of the option. For
         example, '-l __init__ -l 5' will print only the topmost 5 lines of
         information about class constructors.
 
         -r: return the pstats.Stats object generated by the profiling. This
         object has all the information about the profile in it, and you can
         later use it for further analysis or in other functions.
 
        -s <key>: sort profile by given key. You can provide more than one key
         by using the option several times: '-s key1 -s key2 -s key3...'. The
         default sorting key is 'time'.
 
         The following is copied verbatim from the profile documentation
         referenced below:
 
         When more than one key is provided, additional keys are used as
         secondary criteria when the there is equality in all keys selected
         before them.
 
         Abbreviations can be used for any key names, as long as the
         abbreviation is unambiguous.  The following are the keys currently
         defined:
 
                 Valid Arg       Meaning
                   "calls"      call count
                   "cumulative" cumulative time
                   "file"       file name
                   "module"     file name
                   "pcalls"     primitive call count
                   "line"       line number
                   "name"       function name
                   "nfl"        name/file/line
                   "stdname"    standard name
                   "time"       internal time
 
         Note that all sorts on statistics are in descending order (placing
         most time consuming items first), where as name, file, and line number
         searches are in ascending order (i.e., alphabetical). The subtle
         distinction between "nfl" and "stdname" is that the standard name is a
         sort of the name as printed, which means that the embedded line
         numbers get compared in an odd way.  For example, lines 3, 20, and 40
         would (if the file names were the same) appear in the string order
         "20" "3" and "40".  In contrast, "nfl" does a numeric compare of the
         line numbers.  In fact, sort_stats("nfl") is the same as
         sort_stats("name", "file", "line").
 
         -T <filename>: save profile results as shown on screen to a text
         file. The profile is still shown on screen.
 
         -D <filename>: save (via dump_stats) profile statistics to given
         filename. This data is in a format understood by the pstats module, and
         is generated by a call to the dump_stats() method of profile
         objects. The profile is still shown on screen.
 
         -q: suppress output to the pager.  Best used with -T and/or -D above.
 
         If you want to run complete programs under the profiler's control, use
         '%run -p [prof_opts] filename.py [args to program]' where prof_opts
         contains profiler specific options as described here.
 
         You can read the complete documentation for the profile module with::
 
           In [1]: import profile; profile.help()
         """
 
         opts_def = Struct(D=[''],l=[],s=['time'],T=[''])
 
         if user_mode:  # regular user call
             opts,arg_str = self.parse_options(parameter_s,'D:l:rs:T:q',
                                               list_all=True, posix=False)
             namespace = self.shell.user_ns
             if cell is not None:
                 arg_str += '\n' + cell
         else:  # called to run a program by %run -p
             try:
                 filename = get_py_filename(arg_lst[0])
             except IOError as e:
                 try:
                     msg = str(e)
                 except UnicodeError:
                     msg = e.message
                 error(msg)
                 return
 
             arg_str = 'execfile(filename,prog_ns)'
             namespace = {
                 'execfile': self.shell.safe_execfile,
                 'prog_ns': prog_ns,
                 'filename': filename
                 }
 
         opts.merge(opts_def)
 
         prof = profile.Profile()
         try:
             prof = prof.runctx(arg_str,namespace,namespace)
             sys_exit = ''
         except SystemExit:
             sys_exit = """*** SystemExit exception caught in code being profiled."""
 
         stats = pstats.Stats(prof).strip_dirs().sort_stats(*opts.s)
 
         lims = opts.l
         if lims:
             lims = []  # rebuild lims with ints/floats/strings
             for lim in opts.l:
                 try:
                     lims.append(int(lim))
                 except ValueError:
                     try:
                         lims.append(float(lim))
                     except ValueError:
                         lims.append(lim)
 
         # Trap output.
         stdout_trap = StringIO()
 
         if hasattr(stats,'stream'):
             # In newer versions of python, the stats object has a 'stream'
             # attribute to write into.
             stats.stream = stdout_trap
             stats.print_stats(*lims)
         else:
             # For older versions, we manually redirect stdout during printing
             sys_stdout = sys.stdout
             try:
                 sys.stdout = stdout_trap
                 stats.print_stats(*lims)
             finally:
                 sys.stdout = sys_stdout
 
         output = stdout_trap.getvalue()
         output = output.rstrip()
 
         if 'q' not in opts:
             page.page(output)
         print sys_exit,
 
         dump_file = opts.D[0]
         text_file = opts.T[0]
         if dump_file:
             dump_file = unquote_filename(dump_file)
             prof.dump_stats(dump_file)
             print '\n*** Profile stats marshalled to file',\
                   `dump_file`+'.',sys_exit
         if text_file:
             text_file = unquote_filename(text_file)
             pfile = open(text_file,'w')
             pfile.write(output)
             pfile.close()
             print '\n*** Profile printout saved to text file',\
                   `text_file`+'.',sys_exit
 
         if opts.has_key('r'):
             return stats
         else:
             return None
 
     @line_magic
     def pdb(self, parameter_s=''):
         """Control the automatic calling of the pdb interactive debugger.
 
         Call as '%pdb on', '%pdb 1', '%pdb off' or '%pdb 0'. If called without
         argument it works as a toggle.
 
         When an exception is triggered, IPython can optionally call the
         interactive pdb debugger after the traceback printout. %pdb toggles
         this feature on and off.
 
         The initial state of this feature is set in your configuration
         file (the option is ``InteractiveShell.pdb``).
 
         If you want to just activate the debugger AFTER an exception has fired,
         without having to type '%pdb on' and rerunning your code, you can use
         the %debug magic."""
 
         par = parameter_s.strip().lower()
 
         if par:
             try:
                 new_pdb = {'off':0,'0':0,'on':1,'1':1}[par]
             except KeyError:
                 print ('Incorrect argument. Use on/1, off/0, '
                        'or nothing for a toggle.')
                 return
         else:
             # toggle
             new_pdb = not self.shell.call_pdb
 
         # set on the shell
         self.shell.call_pdb = new_pdb
         print 'Automatic pdb calling has been turned',on_off(new_pdb)
 
     @line_magic
     def debug(self, parameter_s=''):
         """Activate the interactive debugger in post-mortem mode.
 
         If an exception has just occurred, this lets you inspect its stack
         frames interactively.  Note that this will always work only on the last
         traceback that occurred, so you must call this quickly after an
         exception that you wish to inspect has fired, because if another one
         occurs, it clobbers the previous one.
 
         If you want IPython to automatically do this on every exception, see
         the %pdb magic for more details.
         """
         self.shell.debugger(force=True)
 
     @line_magic
     def tb(self, s):
         """Print the last traceback with the currently active exception mode.
 
         See %xmode for changing exception reporting modes."""
         self.shell.showtraceback()
 
     @skip_doctest
     @line_magic
     def run(self, parameter_s='', runner=None,
                   file_finder=get_py_filename):
         """Run the named file inside IPython as a program.
 
         Usage:\\
           %run [-n -i -t [-N<N>] -d [-b<N>] -p [profile options]] file [args]
 
         Parameters after the filename are passed as command-line arguments to
         the program (put in sys.argv). Then, control returns to IPython's
         prompt.
 
         This is similar to running at a system prompt:\\
           $ python file args\\
         but with the advantage of giving you IPython's tracebacks, and of
         loading all variables into your interactive namespace for further use
         (unless -p is used, see below).
 
         The file is executed in a namespace initially consisting only of
         __name__=='__main__' and sys.argv constructed as indicated. It thus
         sees its environment as if it were being run as a stand-alone program
         (except for sharing global objects such as previously imported
         modules). But after execution, the IPython interactive namespace gets
         updated with all variables defined in the program (except for __name__
         and sys.argv). This allows for very convenient loading of code for
         interactive work, while giving each program a 'clean sheet' to run in.
 
         Options:
 
         -n: __name__ is NOT set to '__main__', but to the running file's name
         without extension (as python does under import).  This allows running
         scripts and reloading the definitions in them without calling code
         protected by an ' if __name__ == "__main__" ' clause.
 
         -i: run the file in IPython's namespace instead of an empty one. This
         is useful if you are experimenting with code written in a text editor
         which depends on variables defined interactively.
 
         -e: ignore sys.exit() calls or SystemExit exceptions in the script
         being run.  This is particularly useful if IPython is being used to
         run unittests, which always exit with a sys.exit() call.  In such
         cases you are interested in the output of the test results, not in
         seeing a traceback of the unittest module.
 
         -t: print timing information at the end of the run.  IPython will give
         you an estimated CPU time consumption for your script, which under
         Unix uses the resource module to avoid the wraparound problems of
         time.clock().  Under Unix, an estimate of time spent on system tasks
         is also given (for Windows platforms this is reported as 0.0).
 
         If -t is given, an additional -N<N> option can be given, where <N>
         must be an integer indicating how many times you want the script to
         run.  The final timing report will include total and per run results.
 
         For example (testing the script uniq_stable.py)::
 
             In [1]: run -t uniq_stable
 
             IPython CPU timings (estimated):\\
               User  :    0.19597 s.\\
               System:        0.0 s.\\
 
             In [2]: run -t -N5 uniq_stable
 
             IPython CPU timings (estimated):\\
             Total runs performed: 5\\
               Times :      Total       Per run\\
               User  :   0.910862 s,  0.1821724 s.\\
               System:        0.0 s,        0.0 s.
 
         -d: run your program under the control of pdb, the Python debugger.
         This allows you to execute your program step by step, watch variables,
         etc.  Internally, what IPython does is similar to calling:
 
           pdb.run('execfile("YOURFILENAME")')
 
         with a breakpoint set on line 1 of your file.  You can change the line
         number for this automatic breakpoint to be <N> by using the -bN option
         (where N must be an integer).  For example::
 
           %run -d -b40 myscript
 
         will set the first breakpoint at line 40 in myscript.py.  Note that
         the first breakpoint must be set on a line which actually does
         something (not a comment or docstring) for it to stop execution.
 
         When the pdb debugger starts, you will see a (Pdb) prompt.  You must
         first enter 'c' (without quotes) to start execution up to the first
         breakpoint.
 
         Entering 'help' gives information about the use of the debugger.  You
         can easily see pdb's full documentation with "import pdb;pdb.help()"
         at a prompt.
 
         -p: run program under the control of the Python profiler module (which
         prints a detailed report of execution times, function calls, etc).
 
         You can pass other options after -p which affect the behavior of the
         profiler itself. See the docs for %prun for details.
 
         In this mode, the program's variables do NOT propagate back to the
         IPython interactive namespace (because they remain in the namespace
         where the profiler executes them).
 
         Internally this triggers a call to %prun, see its documentation for
         details on the options available specifically for profiling.
 
         There is one special usage for which the text above doesn't apply:
         if the filename ends with .ipy, the file is run as ipython script,
         just as if the commands were written on IPython prompt.
 
         -m: specify module name to load instead of script path. Similar to
         the -m option for the python interpreter. Use this option last if you
         want to combine with other %run options. Unlike the python interpreter
         only source modules are allowed no .pyc or .pyo files.
         For example::
 
             %run -m example
 
         will run the example module.
 
         """
 
         # get arguments and set sys.argv for program to be run.
         opts, arg_lst = self.parse_options(parameter_s, 'nidtN:b:pD:l:rs:T:em:',
                                            mode='list', list_all=1)
         if "m" in opts:
             modulename = opts["m"][0]
             modpath = find_mod(modulename)
             if modpath is None:
                 warn('%r is not a valid modulename on sys.path'%modulename)
                 return
             arg_lst = [modpath] + arg_lst
         try:
             filename = file_finder(arg_lst[0])
         except IndexError:
             warn('you must provide at least a filename.')
             print '\n%run:\n', oinspect.getdoc(self.run)
             return
         except IOError as e:
             try:
                 msg = str(e)
             except UnicodeError:
                 msg = e.message
             error(msg)
             return
 
         if filename.lower().endswith('.ipy'):
             self.shell.safe_execfile_ipy(filename)
             return
 
         # Control the response to exit() calls made by the script being run
         exit_ignore = 'e' in opts
 
         # Make sure that the running script gets a proper sys.argv as if it
         # were run from a system shell.
         save_argv = sys.argv # save it for later restoring
 
         # simulate shell expansion on arguments, at least tilde expansion
         args = [ os.path.expanduser(a) for a in arg_lst[1:] ]
 
         sys.argv = [filename] + args  # put in the proper filename
         # protect sys.argv from potential unicode strings on Python 2:
         if not py3compat.PY3:
             sys.argv = [ py3compat.cast_bytes(a) for a in sys.argv ]
 
         if 'i' in opts:
             # Run in user's interactive namespace
             prog_ns = self.shell.user_ns
             __name__save = self.shell.user_ns['__name__']
             prog_ns['__name__'] = '__main__'
             main_mod = self.shell.new_main_mod(prog_ns)
         else:
             # Run in a fresh, empty namespace
             if 'n' in opts:
                 name = os.path.splitext(os.path.basename(filename))[0]
             else:
                 name = '__main__'
 
             main_mod = self.shell.new_main_mod()
             prog_ns = main_mod.__dict__
             prog_ns['__name__'] = name
 
         # Since '%run foo' emulates 'python foo.py' at the cmd line, we must
         # set the __file__ global in the script's namespace
         prog_ns['__file__'] = filename
 
         # pickle fix.  See interactiveshell for an explanation.  But we need to
         # make sure that, if we overwrite __main__, we replace it at the end
         main_mod_name = prog_ns['__name__']
 
         if main_mod_name == '__main__':
             restore_main = sys.modules['__main__']
         else:
             restore_main = False
 
         # This needs to be undone at the end to prevent holding references to
         # every single object ever created.
         sys.modules[main_mod_name] = main_mod
 
         try:
             stats = None
             with self.shell.readline_no_record:
                 if 'p' in opts:
                     stats = self.prun('', None, False, opts, arg_lst, prog_ns)
                 else:
                     if 'd' in opts:
                         deb = debugger.Pdb(self.shell.colors)
                         # reset Breakpoint state, which is moronically kept
                         # in a class
                         bdb.Breakpoint.next = 1
                         bdb.Breakpoint.bplist = {}
                         bdb.Breakpoint.bpbynumber = [None]
                         # Set an initial breakpoint to stop execution
                         maxtries = 10
                         bp = int(opts.get('b', [1])[0])
                         checkline = deb.checkline(filename, bp)
                         if not checkline:
                             for bp in range(bp + 1, bp + maxtries + 1):
                                 if deb.checkline(filename, bp):
                                     break
                             else:
                                 msg = ("\nI failed to find a valid line to set "
                                        "a breakpoint\n"
                                        "after trying up to line: %s.\n"
                                        "Please set a valid breakpoint manually "
                                        "with the -b option." % bp)
                                 error(msg)
                                 return
                         # if we find a good linenumber, set the breakpoint
                         deb.do_break('%s:%s' % (filename, bp))
                         # Start file run
                         print "NOTE: Enter 'c' at the",
                         print "%s prompt to start your script." % deb.prompt
                         ns = {'execfile': py3compat.execfile, 'prog_ns': prog_ns}
                         try:
                             deb.run('execfile("%s", prog_ns)' % filename, ns)
 
                         except:
                             etype, value, tb = sys.exc_info()
                             # Skip three frames in the traceback: the %run one,
                             # one inside bdb.py, and the command-line typed by the
                             # user (run by exec in pdb itself).
                             self.shell.InteractiveTB(etype, value, tb, tb_offset=3)
                     else:
                         if runner is None:
                             runner = self.default_runner
                         if runner is None:
                             runner = self.shell.safe_execfile
                         if 't' in opts:
                             # timed execution
                             try:
                                 nruns = int(opts['N'][0])
                                 if nruns < 1:
                                     error('Number of runs must be >=1')
                                     return
                             except (KeyError):
                                 nruns = 1
                             twall0 = time.time()
                             if nruns == 1:
                                 t0 = clock2()
                                 runner(filename, prog_ns, prog_ns,
                                        exit_ignore=exit_ignore)
                                 t1 = clock2()
                                 t_usr = t1[0] - t0[0]
                                 t_sys = t1[1] - t0[1]
                                 print "\nIPython CPU timings (estimated):"
                                 print "  User   : %10.2f s." % t_usr
                                 print "  System : %10.2f s." % t_sys
                             else:
                                 runs = range(nruns)
                                 t0 = clock2()
                                 for nr in runs:
                                     runner(filename, prog_ns, prog_ns,
                                            exit_ignore=exit_ignore)
                                 t1 = clock2()
                                 t_usr = t1[0] - t0[0]
                                 t_sys = t1[1] - t0[1]
                                 print "\nIPython CPU timings (estimated):"
                                 print "Total runs performed:", nruns
                                 print "  Times  : %10.2f    %10.2f" % ('Total', 'Per run')
                                 print "  User   : %10.2f s, %10.2f s." % (t_usr, t_usr / nruns)
                                 print "  System : %10.2f s, %10.2f s." % (t_sys, t_sys / nruns)
                             twall1 = time.time()
                             print "Wall time: %10.2f s." % (twall1 - twall0)
 
                         else:
                             # regular execution
                             runner(filename, prog_ns, prog_ns, exit_ignore=exit_ignore)
 
                 if 'i' in opts:
                     self.shell.user_ns['__name__'] = __name__save
                 else:
                     # The shell MUST hold a reference to prog_ns so after %run
                     # exits, the python deletion mechanism doesn't zero it out
                     # (leaving dangling references).
                     self.shell.cache_main_mod(prog_ns, filename)
                     # update IPython interactive namespace
 
                     # Some forms of read errors on the file may mean the
                     # __name__ key was never set; using pop we don't have to
                     # worry about a possible KeyError.
                     prog_ns.pop('__name__', None)
 
                     self.shell.user_ns.update(prog_ns)
         finally:
             # It's a bit of a mystery why, but __builtins__ can change from
             # being a module to becoming a dict missing some key data after
             # %run.  As best I can see, this is NOT something IPython is doing
             # at all, and similar problems have been reported before:
             # http://coding.derkeiler.com/Archive/Python/comp.lang.python/2004-10/0188.html
             # Since this seems to be done by the interpreter itself, the best
             # we can do is to at least restore __builtins__ for the user on
             # exit.
             self.shell.user_ns['__builtins__'] = builtin_mod
 
             # Ensure key global structures are restored
             sys.argv = save_argv
             if restore_main:
                 sys.modules['__main__'] = restore_main
             else:
                 # Remove from sys.modules the reference to main_mod we'd
                 # added.  Otherwise it will trap references to objects
                 # contained therein.
                 del sys.modules[main_mod_name]
 
         return stats
 
     @skip_doctest
     @line_cell_magic
     def timeit(self, line='', cell=None):
         """Time execution of a Python statement or expression
 
         Usage, in line mode:
           %timeit [-n<N> -r<R> [-t|-c]] statement
         or in cell mode:
           %%timeit [-n<N> -r<R> [-t|-c]] setup_code
           code
           code...
 
         Time execution of a Python statement or expression using the timeit
         module.  This function can be used both as a line and cell magic:
 
         - In line mode you can time a single-line statement (though multiple
           ones can be chained with using semicolons).
 
         - In cell mode, the statement in the first line is used as setup code
           (executed but not timed) and the body of the cell is timed.  The cell
           body has access to any variables created in the setup code.
 
         Options:
         -n<N>: execute the given statement <N> times in a loop. If this value
         is not given, a fitting value is chosen.
 
         -r<R>: repeat the loop iteration <R> times and take the best result.
         Default: 3
 
         -t: use time.time to measure the time, which is the default on Unix.
         This function measures wall time.
 
         -c: use time.clock to measure the time, which is the default on
         Windows and measures wall time. On Unix, resource.getrusage is used
         instead and returns the CPU user time.
 
         -p<P>: use a precision of <P> digits to display the timing result.
         Default: 3
 
 
         Examples
         --------
         ::
 
           In [1]: %timeit pass
           10000000 loops, best of 3: 53.3 ns per loop
 
           In [2]: u = None
 
           In [3]: %timeit u is None
           10000000 loops, best of 3: 184 ns per loop
 
           In [4]: %timeit -r 4 u == None
           1000000 loops, best of 4: 242 ns per loop
 
           In [5]: import time
 
           In [6]: %timeit -n1 time.sleep(2)
           1 loops, best of 3: 2 s per loop
 
 
         The times reported by %timeit will be slightly higher than those
         reported by the timeit.py script when variables are accessed. This is
         due to the fact that %timeit executes the statement in the namespace
         of the shell, compared with timeit.py, which uses a single setup
         statement to import function or create variables. Generally, the bias
         does not matter as long as results from timeit.py are not mixed with
         those from %timeit."""
 
         import timeit
         import math
 
         # XXX: Unfortunately the unicode 'micro' symbol can cause problems in
         # certain terminals.  Until we figure out a robust way of
         # auto-detecting if the terminal can deal with it, use plain 'us' for
         # microseconds.  I am really NOT happy about disabling the proper
         # 'micro' prefix, but crashing is worse... If anyone knows what the
         # right solution for this is, I'm all ears...
         #
         # Note: using
         #
         # s = u'\xb5'
         # s.encode(sys.getdefaultencoding())
         #
         # is not sufficient, as I've seen terminals where that fails but
         # print s
         #
         # succeeds
         #
         # See bug: https://bugs.launchpad.net/ipython/+bug/348466
 
         #units = [u"s", u"ms",u'\xb5',"ns"]
         units = [u"s", u"ms",u'us',"ns"]
 
         scaling = [1, 1e3, 1e6, 1e9]
 
         opts, stmt = self.parse_options(line,'n:r:tcp:',
                                         posix=False, strict=False)
         if stmt == "" and cell is None:
             return
         timefunc = timeit.default_timer
         number = int(getattr(opts, "n", 0))
         repeat = int(getattr(opts, "r", timeit.default_repeat))
         precision = int(getattr(opts, "p", 3))
         if hasattr(opts, "t"):
             timefunc = time.time
         if hasattr(opts, "c"):
             timefunc = clock
 
         timer = timeit.Timer(timer=timefunc)
         # this code has tight coupling to the inner workings of timeit.Timer,
         # but is there a better way to achieve that the code stmt has access
         # to the shell namespace?
 
         if cell is None:
             # called as line magic
             setup = 'pass'
             stmt = timeit.reindent(stmt, 8)
         else:
             setup = timeit.reindent(stmt, 4)
             stmt = timeit.reindent(cell, 8)
 
         # From Python 3.3, this template uses new-style string formatting.
         if sys.version_info >= (3, 3):
             src = timeit.template.format(stmt=stmt, setup=setup)
         else:
             src = timeit.template % dict(stmt=stmt, setup=setup)
 
         # Track compilation time so it can be reported if too long
         # Minimum time above which compilation time will be reported
         tc_min = 0.1
 
         t0 = clock()
         code = compile(src, "<magic-timeit>", "exec")
         tc = clock()-t0
 
         ns = {}
         exec code in self.shell.user_ns, ns
         timer.inner = ns["inner"]
 
         if number == 0:
             # determine number so that 0.2 <= total time < 2.0
             number = 1
             for i in range(1, 10):
                 if timer.timeit(number) >= 0.2:
                     break
                 number *= 10
 
         best = min(timer.repeat(repeat, number)) / number
 
         if best > 0.0 and best < 1000.0:
             order = min(-int(math.floor(math.log10(best)) // 3), 3)
         elif best >= 1000.0:
             order = 0
         else:
             order = 3
         print u"%d loops, best of %d: %.*g %s per loop" % (number, repeat,
                                                           precision,
                                                           best * scaling[order],
                                                           units[order])
         if tc > tc_min:
             print "Compiler time: %.2f s" % tc
 
     @skip_doctest
     @needs_local_scope
     @line_magic
     def time(self,parameter_s, user_locals):
         """Time execution of a Python statement or expression.
 
         The CPU and wall clock times are printed, and the value of the
         expression (if any) is returned.  Note that under Win32, system time
         is always reported as 0, since it can not be measured.
 
         This function provides very basic timing functionality.  In Python
         2.3, the timeit module offers more control and sophistication, so this
         could be rewritten to use it (patches welcome).
 
         Examples
         --------
         ::
 
           In [1]: time 2**128
           CPU times: user 0.00 s, sys: 0.00 s, total: 0.00 s
           Wall time: 0.00
           Out[1]: 340282366920938463463374607431768211456L
 
           In [2]: n = 1000000
 
           In [3]: time sum(range(n))
           CPU times: user 1.20 s, sys: 0.05 s, total: 1.25 s
           Wall time: 1.37
           Out[3]: 499999500000L
 
           In [4]: time print 'hello world'
           hello world
           CPU times: user 0.00 s, sys: 0.00 s, total: 0.00 s
           Wall time: 0.00
 
           Note that the time needed by Python to compile the given expression
           will be reported if it is more than 0.1s.  In this example, the
           actual exponentiation is done by Python at compilation time, so while
           the expression can take a noticeable amount of time to compute, that
           time is purely due to the compilation:
 
           In [5]: time 3**9999;
           CPU times: user 0.00 s, sys: 0.00 s, total: 0.00 s
           Wall time: 0.00 s
 
           In [6]: time 3**999999;
           CPU times: user 0.00 s, sys: 0.00 s, total: 0.00 s
           Wall time: 0.00 s
           Compiler : 0.78 s
           """
 
         # fail immediately if the given expression can't be compiled
 
         expr = self.shell.prefilter(parameter_s,False)
 
         # Minimum time above which compilation time will be reported
         tc_min = 0.1
 
         try:
             mode = 'eval'
             t0 = clock()
             code = compile(expr,'<timed eval>',mode)
             tc = clock()-t0
         except SyntaxError:
             mode = 'exec'
             t0 = clock()
             code = compile(expr,'<timed exec>',mode)
             tc = clock()-t0
         # skew measurement as little as possible
         glob = self.shell.user_ns
         wtime = time.time
         # time execution
         wall_st = wtime()
         if mode=='eval':
             st = clock2()
             out = eval(code, glob, user_locals)
             end = clock2()
         else:
             st = clock2()
             exec code in glob, user_locals
             end = clock2()
             out = None
         wall_end = wtime()
         # Compute actual times and report
         wall_time = wall_end-wall_st
         cpu_user = end[0]-st[0]
         cpu_sys = end[1]-st[1]
         cpu_tot = cpu_user+cpu_sys
         print "CPU times: user %.2f s, sys: %.2f s, total: %.2f s" % \
               (cpu_user,cpu_sys,cpu_tot)
         print "Wall time: %.2f s" % wall_time
         if tc > tc_min:
             print "Compiler : %.2f s" % tc
         return out
 
     @skip_doctest
     @line_magic
     def macro(self, parameter_s=''):
         """Define a macro for future re-execution. It accepts ranges of history,
         filenames or string objects.
 
         Usage:\\
           %macro [options] name n1-n2 n3-n4 ... n5 .. n6 ...
 
         Options:
 
           -r: use 'raw' input.  By default, the 'processed' history is used,
           so that magics are loaded in their transformed version to valid
           Python.  If this option is given, the raw input as typed as the
           command line is used instead.
 
         This will define a global variable called `name` which is a string
         made of joining the slices and lines you specify (n1,n2,... numbers
         above) from your input history into a single string. This variable
         acts like an automatic function which re-executes those lines as if
         you had typed them. You just type 'name' at the prompt and the code
         executes.
 
         The syntax for indicating input ranges is described in %history.
 
         Note: as a 'hidden' feature, you can also use traditional python slice
         notation, where N:M means numbers N through M-1.
 
         For example, if your history contains (%hist prints it)::
 
           44: x=1
           45: y=3
           46: z=x+y
           47: print x
           48: a=5
           49: print 'x',x,'y',y
 
         you can create a macro with lines 44 through 47 (included) and line 49
         called my_macro with::
 
           In [55]: %macro my_macro 44-47 49
 
         Now, typing `my_macro` (without quotes) will re-execute all this code
         in one pass.
 
         You don't need to give the line-numbers in order, and any given line
         number can appear multiple times. You can assemble macros with any
         lines from your input history in any order.
 
         The macro is a simple object which holds its value in an attribute,
         but IPython's display system checks for macros and executes them as
         code instead of printing them when you type their name.
 
         You can view a macro's contents by explicitly printing it with::
 
           print macro_name
 
         """
         opts,args = self.parse_options(parameter_s,'r',mode='list')
         if not args:   # List existing macros
             return sorted(k for k,v in self.shell.user_ns.iteritems() if\
                                                         isinstance(v, Macro))
         if len(args) == 1:
             raise UsageError(
                 "%macro insufficient args; usage '%macro name n1-n2 n3-4...")
         name, codefrom = args[0], " ".join(args[1:])
 
         #print 'rng',ranges  # dbg
         try:
             lines = self.shell.find_user_code(codefrom, 'r' in opts)
         except (ValueError, TypeError) as e:
             print e.args[0]
             return
         macro = Macro(lines)
         self.shell.define_macro(name, macro)
         print 'Macro `%s` created. To execute, type its name (without quotes).' % name
         print '=== Macro contents: ==='
         print macro,
+    
+    @magic_arguments.magic_arguments()
+    @magic_arguments.argument('output', type=str, default='', nargs='?',
+        help="""The name of the variable in which to store output.
+        This is a utils.io.CapturedIO object with stdout/err attributes
+        for the text of the captured output.
+        
+        CapturedOutput also has a show() method for displaying the output,
+        and __call__ as well, so you can use that to quickly display the
+        output.
+        
+        If unspecified, captured output is discarded.
+        """
+    )
+    @magic_arguments.argument('--no-stderr', action="store_true",
+        help="""Don't capture stderr."""
+    )
+    @magic_arguments.argument('--no-stdout', action="store_true",
+        help="""Don't capture stdout."""
+    )
+    @cell_magic
+    def capture(self, line, cell):
+        """run the cell, capturing stdout/err"""
+        args = magic_arguments.parse_argstring(self.capture, line)
+        out = not args.no_stdout
+        err = not args.no_stderr
+        with capture_output(out, err) as io:
+            self.shell.run_cell(cell)
+        if args.output:
+            self.shell.user_ns[args.output] = 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,384 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 = stdout
+        self._stderr = stderr
+    
+    def __str__(self):
+        return self.stdout
+    
+    @property
+    def stdout(self):
+        if not self._stdout:
+            return ''
+        return self._stdout.getvalue()
+    
+    @property
+    def stderr(self):
+        if not self._stderr:
+            return ''
+        return self._stderr.getvalue()
+    
+    def show(self):
+        """write my output to sys.stdout/err as appropriate"""
+        sys.stdout.write(self.stdout)
+        sys.stderr.write(self.stderr)
+        sys.stdout.flush()
+        sys.stderr.flush()
+    
+    __call__ = show
+
+
+class capture_output(object):
+    """context manager for capturing stdout/err"""
+    stdout = True
+    stderr = True
+    
+    def __init__(self, stdout=True, stderr=True):
+        self.stdout = stdout
+        self.stderr = stderr
+    
+    def __enter__(self):
+        self.sys_stdout = sys.stdout
+        self.sys_stderr = sys.stderr
+        
+        stdout = stderr = False
+        if self.stdout:
+            stdout = sys.stdout = StringIO()
+        if self.stderr:
+            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')