capture.py
166 lines
| 4.9 KiB
| text/x-python
|
PythonLexer
|
r12223 | # encoding: utf-8 | ||
|
r16585 | """IO capturing utilities.""" | ||
|
r12223 | |||
|
r16585 | # Copyright (c) IPython Development Team. | ||
# Distributed under the terms of the Modified BSD License. | ||||
|
r12223 | |||
import sys | ||||
|
r23111 | from io import StringIO | ||
|
r12223 | |||
#----------------------------------------------------------------------------- | ||||
# Classes and functions | ||||
#----------------------------------------------------------------------------- | ||||
class RichOutput(object): | ||||
|
r16585 | def __init__(self, data=None, metadata=None): | ||
|
r12223 | self.data = data or {} | ||
self.metadata = metadata or {} | ||||
def display(self): | ||||
from IPython.display import publish_display_data | ||||
|
r16585 | publish_display_data(data=self.data, metadata=self.metadata) | ||
|
r12223 | |||
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): | ||||
|
r12227 | return self._repr_mime_("image/jpeg") | ||
|
r12223 | |||
def _repr_svg_(self): | ||||
return self._repr_mime_("image/svg+xml") | ||||
class CapturedIO(object): | ||||
|
r12225 | """Simple object for containing captured stdout/err and rich display StringIO objects | ||
Each instance `c` has three attributes: | ||||
|
r12553 | - ``c.stdout`` : standard output as a string | ||
- ``c.stderr`` : standard error as a string | ||||
- ``c.outputs``: a list of rich display outputs | ||||
|
r12225 | |||
|
r12553 | Additionally, there's a ``c.show()`` method which will print all of the | ||
above in the same order, and can be invoked simply via ``c()``. | ||||
|
r12225 | """ | ||
|
r12223 | |||
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): | ||||
|
r12226 | "Captured standard output" | ||
|
r12223 | if not self._stdout: | ||
return '' | ||||
return self._stdout.getvalue() | ||||
@property | ||||
def stderr(self): | ||||
|
r12226 | "Captured standard error" | ||
|
r12223 | if not self._stderr: | ||
return '' | ||||
return self._stderr.getvalue() | ||||
|
r12224 | @property | ||
def outputs(self): | ||||
|
r12226 | """A list of the captured rich display outputs, if any. | ||
|
r12553 | If you have a CapturedIO object ``c``, these can be displayed in IPython | ||
using:: | ||||
|
r12226 | |||
from IPython.display import display | ||||
for o in c.outputs: | ||||
display(o) | ||||
""" | ||||
|
r16585 | return [ RichOutput(d, md) for d, md in self._outputs ] | ||
|
r12224 | |||
|
r12223 | 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() | ||||
|
r16585 | for data, metadata in self._outputs: | ||
RichOutput(data, metadata).display() | ||||
|
r12223 | |||
__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 | ||||
|
r22774 | from IPython.core.displayhook import CapturingDisplayHook | ||
|
r12223 | |||
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 | ||||
|
r12225 | stdout = stderr = outputs = None | ||
|
r12223 | 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 | ||||
|
r22774 | self.save_display_hook = sys.displayhook | ||
sys.displayhook = CapturingDisplayHook(shell=self.shell, | ||||
outputs=outputs) | ||||
|
r12223 | |||
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 | ||||
|
r22774 | sys.displayhook = self.save_display_hook | ||
|
r12223 | |||