update capture per review

minrk · minrk · commit 34ddebb4e9af · 2012-06-10T17:34:08.000-07:00
* capture stores io object, rather than stdout/err text directly
* add options to selectively disable capturing stdout/err
diff --git a/IPython/core/magics/execution.py b/IPython/core/magics/execution.py
@@ -992,25 +992,31 @@ def macro(self, parameter_s=''):
         print macro,
     
     @magic_arguments.magic_arguments()
-    @magic_arguments.argument('-o', '--out', type=str,
-        help="""The name of the variable in which to store stdout
+    @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.
         
-        If unspecified: stdout is discarded
-        """
-    )
-    @magic_arguments.argument('-e', '--err', type=str,
-        help="""The name of the variable in which to store stderr
+        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: stderr is discarded
+        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)
-        with capture_output() as io:
+        out = not args.no_stdout
+        err = not args.no_stderr
+        with capture_output(out, err) as io:
             self.shell.run_cell(cell)
-        if args.out:
-            self.shell.user_ns[args.out] = io.stdout
-        if args.err:
-            self.shell.user_ns[args.err] = io.stderr
+        if args.output:
+            self.shell.user_ns[args.output] = io
diff --git a/IPython/utils/io.py b/IPython/utils/io.py
@@ -328,26 +328,50 @@ 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
+        self._stdout = stdout
+        self._stderr = stderr
     
     @property
     def stdout(self):
-        return self.stdout_io.getvalue()
+        if not self._stdout:
+            return ''
+        return self._stdout.getvalue()
     
     @property
     def stderr(self):
-        return self.stderr_io.getvalue()
+        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 = sys.stdout = StringIO()
-        stderr = sys.stderr = StringIO()
+        
+        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):
diff --git a/docs/examples/notebooks/Capturing Output.ipynb b/docs/examples/notebooks/Capturing Output.ipynb
@@ -47,13 +47,13 @@
     {
      "cell_type": "markdown",
      "source": [
-      "If you specify `-o` or `-e`, then stdout and/or stderr will be stored in those variables in your namespace."
+      "If you specify a name, then stdout and stderr will be stored in an object in your namespace."
      ]
     },
     {
      "cell_type": "code",
      "input": [
-      "%%capture -o my_stdout",
+      "%%capture captured",
       "print 'hi, stdout'",
       "print >> sys.stderr, 'hi, stderr'"
      ],
@@ -63,26 +63,37 @@
     {
      "cell_type": "code",
      "input": [
-      "my_stdout"
+      "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": [
-      "%%capture -o my_stdout2 -e my_stderr",
-      "print 'hi again, stdout'",
-      "print >> sys.stderr, 'hi there, stderr'"
+      "captured.stdout"
      ],
      "language": "python",
      "outputs": []
     },
     {
      "cell_type": "code",
      "input": [
-      "sys.stdout.write(my_stdout2)",
-      "sys.stderr.write(my_stderr)"
+      "captured.stderr"
      ],
      "language": "python",
      "outputs": []
@@ -104,7 +115,7 @@
     {
      "cell_type": "code",
      "input": [
-      "%%capture -o wontshutup",
+      "%%capture wontshutup",
       "",
       "print \"setting up X\"",
       "x = np.linspace(0,5,1000)",
@@ -120,7 +131,47 @@
     {
      "cell_type": "code",
      "input": [
-      "print wontshutup"
+      "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": []
