From e10d2ee49916a2b47d563e39f635490d49f9dd12 2013-08-17 02:37:44 From: MinRK Date: 2013-08-17 02:37:44 Subject: [PATCH] capture rich output as well as stdout/err in capture_output closes #3742 --- diff --git a/IPython/core/displaypub.py b/IPython/core/displaypub.py index c12d873..254a5d5 100644 --- a/IPython/core/displaypub.py +++ b/IPython/core/displaypub.py @@ -31,6 +31,7 @@ from __future__ import print_function from IPython.config.configurable import Configurable from IPython.utils import io +from IPython.utils.traitlets import List #----------------------------------------------------------------------------- # Main payload class @@ -115,7 +116,20 @@ class DisplayPublisher(Configurable): if stderr: print('\033[2K\r', file=io.stderr, end='') io.stderr.flush() - + + +class CapturingDisplayPublisher(DisplayPublisher): + """A DisplayPublisher that stores""" + outputs = List() + + def publish(self, source, data, metadata=None): + self.outputs.append((source, data, metadata)) + + def clear_output(self, stdout=True, stderr=True, other=True): + super(CapturingDisplayPublisher, self).clear_output(stdout, stderr, other) + if other: + # empty the list, *do not* reassign a new list + del self.outputs[:] def publish_display_data(source, data, metadata=None): diff --git a/IPython/utils/capture.py b/IPython/utils/capture.py new file mode 100644 index 0000000..b54c120 --- /dev/null +++ b/IPython/utils/capture.py @@ -0,0 +1,148 @@ +# encoding: utf-8 +""" +IO capturing utilities. +""" + +#----------------------------------------------------------------------------- +# Copyright (C) 2013 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 sys +from StringIO import StringIO + +#----------------------------------------------------------------------------- +# Classes and functions +#----------------------------------------------------------------------------- + + +class RichOutput(object): + def __init__(self, source, data, metadata): + self.source = source + self.data = data or {} + self.metadata = metadata or {} + + def display(self): + from IPython.display import publish_display_data + publish_display_data(self.source, self.data, self.metadata) + + def _repr_mime_(self, mime): + if mime not in self.data: + return + data = self.data[mime] + if mime in self.metadata: + return data, self.metadata[mime] + else: + return data + + def _repr_html_(self): + return self._repr_mime_("text/html") + + def _repr_latex_(self): + return self._repr_mime_("text/latex") + + def _repr_json_(self): + return self._repr_mime_("application/json") + + def _repr_javascript_(self): + return self._repr_mime_("application/javascript") + + def _repr_png_(self): + return self._repr_mime_("image/png") + + def _repr_jpeg_(self): + return self._repr_mime_("image/jpg") + + def _repr_svg_(self): + return self._repr_mime_("image/svg+xml") + + +class CapturedIO(object): + """Simple object for containing captured stdout/err StringIO objects""" + + def __init__(self, stdout, stderr, outputs=None): + self._stdout = stdout + self._stderr = stderr + if outputs is None: + outputs = [] + self._outputs = outputs + + 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() + for source, data, metadata in self._outputs: + RichOutput(source, data, metadata).display() + + __call__ = show + + +class capture_output(object): + """context manager for capturing stdout/err""" + stdout = True + stderr = True + display = True + + def __init__(self, stdout=True, stderr=True, display=True): + self.stdout = stdout + self.stderr = stderr + self.display = display + self.shell = None + + def __enter__(self): + from IPython.core.getipython import get_ipython + from IPython.core.displaypub import CapturingDisplayPublisher + + self.sys_stdout = sys.stdout + self.sys_stderr = sys.stderr + + if self.display: + self.shell = get_ipython() + if self.shell is None: + self.save_display_pub = None + self.display = False + + stdout = stderr = outputs = False + if self.stdout: + stdout = sys.stdout = StringIO() + if self.stderr: + stderr = sys.stderr = StringIO() + if self.display: + self.save_display_pub = self.shell.display_pub + self.shell.display_pub = CapturingDisplayPublisher() + outputs = self.shell.display_pub.outputs + + + return CapturedIO(stdout, stderr, outputs) + + def __exit__(self, exc_type, exc_value, traceback): + sys.stdout = self.sys_stdout + sys.stderr = self.sys_stderr + if self.display and self.shell: + self.shell.display_pub = self.save_display_pub + + diff --git a/IPython/utils/io.py b/IPython/utils/io.py index e5be667..ddd21ab 100644 --- a/IPython/utils/io.py +++ b/IPython/utils/io.py @@ -18,6 +18,7 @@ import os import sys import tempfile from StringIO import StringIO +from .capture import CapturedIO, capture_output #----------------------------------------------------------------------------- # Code @@ -226,63 +227,3 @@ 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 - -