Skip to content

Commit e10d2ee

Browse files
minrkminrk
committed
capture rich output as well as stdout/err in capture_output
closes #3742
1 parent e761579 commit e10d2ee

3 files changed

Lines changed: 164 additions & 61 deletions

File tree

‎IPython/core/displaypub.py‎

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131

3232
from IPython.config.configurable import Configurable
3333
from IPython.utils import io
34+
from IPython.utils.traitlets import List
3435

3536
#-----------------------------------------------------------------------------
3637
# Main payload class
@@ -115,7 +116,20 @@ def clear_output(self, stdout=True, stderr=True, other=True):
115116
if stderr:
116117
print('\033[2K\r', file=io.stderr, end='')
117118
io.stderr.flush()
118-
119+
120+
121+
class CapturingDisplayPublisher(DisplayPublisher):
122+
"""A DisplayPublisher that stores"""
123+
outputs = List()
124+
125+
def publish(self, source, data, metadata=None):
126+
self.outputs.append((source, data, metadata))
127+
128+
def clear_output(self, stdout=True, stderr=True, other=True):
129+
super(CapturingDisplayPublisher, self).clear_output(stdout, stderr, other)
130+
if other:
131+
# empty the list, *do not* reassign a new list
132+
del self.outputs[:]
119133

120134

121135
def publish_display_data(source, data, metadata=None):

‎IPython/utils/capture.py‎

Lines changed: 148 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,148 @@
1+
# encoding: utf-8
2+
"""
3+
IO capturing utilities.
4+
"""
5+
6+
#-----------------------------------------------------------------------------
7+
# Copyright (C) 2013 The IPython Development Team
8+
#
9+
# Distributed under the terms of the BSD License. The full license is in
10+
# the file COPYING, distributed as part of this software.
11+
#-----------------------------------------------------------------------------
12+
from __future__ import print_function
13+
14+
#-----------------------------------------------------------------------------
15+
# Imports
16+
#-----------------------------------------------------------------------------
17+
18+
import sys
19+
from StringIO import StringIO
20+
21+
#-----------------------------------------------------------------------------
22+
# Classes and functions
23+
#-----------------------------------------------------------------------------
24+
25+
26+
class RichOutput(object):
27+
def __init__(self, source, data, metadata):
28+
self.source = source
29+
self.data = data or {}
30+
self.metadata = metadata or {}
31+
32+
def display(self):
33+
from IPython.display import publish_display_data
34+
publish_display_data(self.source, self.data, self.metadata)
35+
36+
def _repr_mime_(self, mime):
37+
if mime not in self.data:
38+
return
39+
data = self.data[mime]
40+
if mime in self.metadata:
41+
return data, self.metadata[mime]
42+
else:
43+
return data
44+
45+
def _repr_html_(self):
46+
return self._repr_mime_("text/html")
47+
48+
def _repr_latex_(self):
49+
return self._repr_mime_("text/latex")
50+
51+
def _repr_json_(self):
52+
return self._repr_mime_("application/json")
53+
54+
def _repr_javascript_(self):
55+
return self._repr_mime_("application/javascript")
56+
57+
def _repr_png_(self):
58+
return self._repr_mime_("image/png")
59+
60+
def _repr_jpeg_(self):
61+
return self._repr_mime_("image/jpg")
62+
63+
def _repr_svg_(self):
64+
return self._repr_mime_("image/svg+xml")
65+
66+
67+
class CapturedIO(object):
68+
"""Simple object for containing captured stdout/err StringIO objects"""
69+
70+
def __init__(self, stdout, stderr, outputs=None):
71+
self._stdout = stdout
72+
self._stderr = stderr
73+
if outputs is None:
74+
outputs = []
75+
self._outputs = outputs
76+
77+
def __str__(self):
78+
return self.stdout
79+
80+
@property
81+
def stdout(self):
82+
if not self._stdout:
83+
return ''
84+
return self._stdout.getvalue()
85+
86+
@property
87+
def stderr(self):
88+
if not self._stderr:
89+
return ''
90+
return self._stderr.getvalue()
91+
92+
def show(self):
93+
"""write my output to sys.stdout/err as appropriate"""
94+
sys.stdout.write(self.stdout)
95+
sys.stderr.write(self.stderr)
96+
sys.stdout.flush()
97+
sys.stderr.flush()
98+
for source, data, metadata in self._outputs:
99+
RichOutput(source, data, metadata).display()
100+
101+
__call__ = show
102+
103+
104+
class capture_output(object):
105+
"""context manager for capturing stdout/err"""
106+
stdout = True
107+
stderr = True
108+
display = True
109+
110+
def __init__(self, stdout=True, stderr=True, display=True):
111+
self.stdout = stdout
112+
self.stderr = stderr
113+
self.display = display
114+
self.shell = None
115+
116+
def __enter__(self):
117+
from IPython.core.getipython import get_ipython
118+
from IPython.core.displaypub import CapturingDisplayPublisher
119+
120+
self.sys_stdout = sys.stdout
121+
self.sys_stderr = sys.stderr
122+
123+
if self.display:
124+
self.shell = get_ipython()
125+
if self.shell is None:
126+
self.save_display_pub = None
127+
self.display = False
128+
129+
stdout = stderr = outputs = False
130+
if self.stdout:
131+
stdout = sys.stdout = StringIO()
132+
if self.stderr:
133+
stderr = sys.stderr = StringIO()
134+
if self.display:
135+
self.save_display_pub = self.shell.display_pub
136+
self.shell.display_pub = CapturingDisplayPublisher()
137+
outputs = self.shell.display_pub.outputs
138+
139+
140+
return CapturedIO(stdout, stderr, outputs)
141+
142+
def __exit__(self, exc_type, exc_value, traceback):
143+
sys.stdout = self.sys_stdout
144+
sys.stderr = self.sys_stderr
145+
if self.display and self.shell:
146+
self.shell.display_pub = self.save_display_pub
147+
148+

‎IPython/utils/io.py‎

Lines changed: 1 addition & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import sys
1919
import tempfile
2020
from StringIO import StringIO
21+
from .capture import CapturedIO, capture_output
2122

2223
#-----------------------------------------------------------------------------
2324
# Code
@@ -226,63 +227,3 @@ def raw_print_err(*args, **kw):
226227
# Short aliases for quick debugging, do NOT use these in production code.
227228
rprint = raw_print
228229
rprinte = raw_print_err
229-
230-
231-
class CapturedIO(object):
232-
"""Simple object for containing captured stdout/err StringIO objects"""
233-
234-
def __init__(self, stdout, stderr):
235-
self._stdout = stdout
236-
self._stderr = stderr
237-
238-
def __str__(self):
239-
return self.stdout
240-
241-
@property
242-
def stdout(self):
243-
if not self._stdout:
244-
return ''
245-
return self._stdout.getvalue()
246-
247-
@property
248-
def stderr(self):
249-
if not self._stderr:
250-
return ''
251-
return self._stderr.getvalue()
252-
253-
def show(self):
254-
"""write my output to sys.stdout/err as appropriate"""
255-
sys.stdout.write(self.stdout)
256-
sys.stderr.write(self.stderr)
257-
sys.stdout.flush()
258-
sys.stderr.flush()
259-
260-
__call__ = show
261-
262-
263-
class capture_output(object):
264-
"""context manager for capturing stdout/err"""
265-
stdout = True
266-
stderr = True
267-
268-
def __init__(self, stdout=True, stderr=True):
269-
self.stdout = stdout
270-
self.stderr = stderr
271-
272-
def __enter__(self):
273-
self.sys_stdout = sys.stdout
274-
self.sys_stderr = sys.stderr
275-
276-
stdout = stderr = False
277-
if self.stdout:
278-
stdout = sys.stdout = StringIO()
279-
if self.stderr:
280-
stderr = sys.stderr = StringIO()
281-
282-
return CapturedIO(stdout, stderr)
283-
284-
def __exit__(self, exc_type, exc_value, traceback):
285-
sys.stdout = self.sys_stdout
286-
sys.stderr = self.sys_stderr
287-
288-

0 commit comments

Comments
 (0)