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