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