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

@@ -33,13 +33,15 b' except ImportError:'
 
 # 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
@@ -988,3 +990,33 b' python-profiler package from non-free.""")'
         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

@@ -100,36 +100,6 b' def skip_without(*names):'
 # Classes
 #-------------------------------------------------------------------------------
 
-class CapturedIO(object):
-    """Simple object for containing captured stdout/err StringIO objects"""
-    
-    def __init__(self, stdout, stderr):
-        self.stdout_io = stdout
-        self.stderr_io = stderr
-    
-    @property
-    def stdout(self):
-        return self.stdout_io.getvalue()
-    
-    @property
-    def stderr(self):
-        return self.stderr_io.getvalue()
-
-
-class capture_output(object):
-    """context manager for capturing stdout/err"""
-    
-    def __enter__(self):
-        self.sys_stdout = sys.stdout
-        self.sys_stderr = sys.stderr
-        stdout = sys.stdout = StringIO()
-        stderr = sys.stderr = StringIO()
-        return CapturedIO(stdout, stderr)
-    
-    def __exit__(self, exc_type, exc_value, traceback):
-        sys.stdout = self.sys_stdout
-        sys.stderr = self.sys_stderr
-
 
 class ClusterTestCase(BaseZMQTestCase):
     

IPython/parallel/tests/test_asyncresult.py

@@ -18,11 +18,12 b' Authors:'
 
 import time
 
-from IPython.parallel.error import TimeoutError
+from IPython.utils.io import capture_output
 
+from IPython.parallel.error import TimeoutError
 from IPython.parallel import error, Client
 from IPython.parallel.tests import add_engines
-from .clienttest import ClusterTestCase, capture_output
+from .clienttest import ClusterTestCase
 
 def setup():
     add_engines(2, total=True)

IPython/parallel/tests/test_magics.py

@@ -25,6 +25,7 b' from nose import SkipTest'
 
 from IPython.testing import decorators as dec
 from IPython.testing.ipunittest import ParametricTestCase
+from IPython.utils.io import capture_output
 
 from IPython import parallel  as pmod
 from IPython.parallel import error
@@ -33,7 +34,7 b' 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)

IPython/utils/io.py

@@ -17,6 +17,7 b' from __future__ import print_function'
 import os
 import sys
 import tempfile
+from StringIO import StringIO
 
 #-----------------------------------------------------------------------------
 # Code
@@ -321,3 +322,63 b' def raw_print_err(*args, **kw):'
 # Short aliases for quick debugging, do NOT use these in production code.
 rprint = raw_print
 rprinte = raw_print_err
+
+
+class CapturedIO(object):
+    """Simple object for containing captured stdout/err StringIO objects"""
+    
+    def __init__(self, stdout, stderr):
+        self._stdout = 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

@@ -20,7 +20,7 b' from subprocess import Popen, PIPE'
 import nose.tools as nt
 
 from IPython.testing import decorators as dec
-from IPython.utils.io import Tee
+from IPython.utils.io import Tee, capture_output
 from IPython.utils.py3compat import doctest_refactor_print
 
 #-----------------------------------------------------------------------------
@@ -73,3 +73,13 b' def test_io_init():'
         # __class__ is a reference to the class object in Python 3, so we can't
         # just test for string equality.
         assert 'IPython.utils.io.IOStream' in classname, classname
+
+def test_capture_output():
+    """capture_output() context works"""
+    
+    with capture_output() as io:
+        print 'hi, stdout'
+        print >> sys.stderr, 'hi, stderr'
+    
+    nt.assert_equals(io.stdout, 'hi, stdout\n')
+    nt.assert_equals(io.stderr, 'hi, stderr\n')