diff --git a/IPython/core/magics/execution.py b/IPython/core/magics/execution.py index 3ec28e1..439792b 100644 --- a/IPython/core/magics/execution.py +++ b/IPython/core/magics/execution.py @@ -33,13 +33,15 @@ 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 @@ 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 diff --git a/IPython/parallel/tests/clienttest.py b/IPython/parallel/tests/clienttest.py index 60017b6..478304d 100644 --- a/IPython/parallel/tests/clienttest.py +++ b/IPython/parallel/tests/clienttest.py @@ -100,36 +100,6 @@ def skip_without(*names): # Classes #------------------------------------------------------------------------------- -class CapturedIO(object): - """Simple object for containing captured stdout/err StringIO objects""" - - def __init__(self, stdout, stderr): - self.stdout_io = stdout - self.stderr_io = stderr - - @property - def stdout(self): - return self.stdout_io.getvalue() - - @property - def stderr(self): - return self.stderr_io.getvalue() - - -class capture_output(object): - """context manager for capturing stdout/err""" - - def __enter__(self): - self.sys_stdout = sys.stdout - self.sys_stderr = sys.stderr - stdout = sys.stdout = StringIO() - stderr = sys.stderr = StringIO() - return CapturedIO(stdout, stderr) - - def __exit__(self, exc_type, exc_value, traceback): - sys.stdout = self.sys_stdout - sys.stderr = self.sys_stderr - class ClusterTestCase(BaseZMQTestCase): diff --git a/IPython/parallel/tests/test_asyncresult.py b/IPython/parallel/tests/test_asyncresult.py index 9e0ba9c..e12f72c 100644 --- a/IPython/parallel/tests/test_asyncresult.py +++ b/IPython/parallel/tests/test_asyncresult.py @@ -18,11 +18,12 @@ Authors: import time -from IPython.parallel.error import TimeoutError +from IPython.utils.io import capture_output +from IPython.parallel.error import TimeoutError from IPython.parallel import error, Client from IPython.parallel.tests import add_engines -from .clienttest import ClusterTestCase, capture_output +from .clienttest import ClusterTestCase def setup(): add_engines(2, total=True) diff --git a/IPython/parallel/tests/test_magics.py b/IPython/parallel/tests/test_magics.py index 7279067..eeda8bd 100644 --- a/IPython/parallel/tests/test_magics.py +++ b/IPython/parallel/tests/test_magics.py @@ -25,6 +25,7 @@ from nose import SkipTest from IPython.testing import decorators as dec from IPython.testing.ipunittest import ParametricTestCase +from IPython.utils.io import capture_output from IPython import parallel as pmod from IPython.parallel import error @@ -33,7 +34,7 @@ from IPython.parallel.util import interactive from IPython.parallel.tests import add_engines -from .clienttest import ClusterTestCase, capture_output, generate_output +from .clienttest import ClusterTestCase, generate_output def setup(): add_engines(3, total=True) diff --git a/IPython/utils/io.py b/IPython/utils/io.py index 3869635..ec00a29 100644 --- a/IPython/utils/io.py +++ b/IPython/utils/io.py @@ -17,6 +17,7 @@ from __future__ import print_function import os import sys import tempfile +from StringIO import StringIO #----------------------------------------------------------------------------- # Code @@ -321,3 +322,63 @@ 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 + + diff --git a/IPython/utils/tests/test_io.py b/IPython/utils/tests/test_io.py index eb49e40..71244d3 100644 --- a/IPython/utils/tests/test_io.py +++ b/IPython/utils/tests/test_io.py @@ -20,7 +20,7 @@ from subprocess import Popen, PIPE import nose.tools as nt from IPython.testing import decorators as dec -from IPython.utils.io import Tee +from IPython.utils.io import Tee, capture_output from IPython.utils.py3compat import doctest_refactor_print #----------------------------------------------------------------------------- @@ -73,3 +73,13 @@ def test_io_init(): # __class__ is a reference to the class object in Python 3, so we can't # just test for string equality. assert 'IPython.utils.io.IOStream' in classname, classname + +def test_capture_output(): + """capture_output() context works""" + + with capture_output() as io: + print 'hi, stdout' + print >> sys.stderr, 'hi, stderr' + + nt.assert_equals(io.stdout, 'hi, stdout\n') + nt.assert_equals(io.stderr, 'hi, stderr\n') diff --git a/docs/examples/notebooks/Capturing Output.ipynb b/docs/examples/notebooks/Capturing Output.ipynb new file mode 100644 index 0000000..eadf591 --- /dev/null +++ b/docs/examples/notebooks/Capturing Output.ipynb @@ -0,0 +1,182 @@ +{ + "metadata": { + "name": "Capturing Output" + }, + "nbformat": 3, + "worksheets": [ + { + "cells": [ + { + "cell_type": "heading", + "level": 1, + "source": [ + "Capturing Output with %%capture" + ] + }, + { + "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