|
| 1 | +# encoding: utf-8 |
| 2 | +"""IO capturing utilities.""" |
| 3 | + |
| 4 | +# Copyright (c) IPython Development Team. |
| 5 | +# Distributed under the terms of the Modified BSD License. |
| 6 | + |
| 7 | +from IPython.core import magic_arguments |
| 8 | +from IPython.core.magic import ( |
| 9 | + Magics, |
| 10 | + cell_magic, |
| 11 | + line_cell_magic, |
| 12 | + line_magic, |
| 13 | + magics_class, |
| 14 | + needs_local_scope, |
| 15 | + no_var_expand, |
| 16 | + on_off, |
| 17 | +) |
| 18 | +from IPython.core.displayhook import DisplayHook |
| 19 | + |
| 20 | + |
| 21 | +import sys |
| 22 | +from io import StringIO |
| 23 | + |
| 24 | +#----------------------------------------------------------------------------- |
| 25 | +# Classes and functions |
| 26 | +#----------------------------------------------------------------------------- |
| 27 | + |
| 28 | + |
| 29 | +class RichOutput(object): |
| 30 | + def __init__(self, data=None, metadata=None, transient=None, update=False): |
| 31 | + self.data = data or {} |
| 32 | + self.metadata = metadata or {} |
| 33 | + self.transient = transient or {} |
| 34 | + self.update = update |
| 35 | + |
| 36 | + def display(self): |
| 37 | + from IPython.display import publish_display_data |
| 38 | + publish_display_data(data=self.data, metadata=self.metadata, |
| 39 | + transient=self.transient, update=self.update) |
| 40 | + |
| 41 | + def _repr_mime_(self, mime): |
| 42 | + if mime not in self.data: |
| 43 | + return |
| 44 | + data = self.data[mime] |
| 45 | + if mime in self.metadata: |
| 46 | + return data, self.metadata[mime] |
| 47 | + else: |
| 48 | + return data |
| 49 | + |
| 50 | + def _repr_mimebundle_(self, include=None, exclude=None): |
| 51 | + return self.data, self.metadata |
| 52 | + |
| 53 | + def _repr_html_(self): |
| 54 | + return self._repr_mime_("text/html") |
| 55 | + |
| 56 | + def _repr_latex_(self): |
| 57 | + return self._repr_mime_("text/latex") |
| 58 | + |
| 59 | + def _repr_json_(self): |
| 60 | + return self._repr_mime_("application/json") |
| 61 | + |
| 62 | + def _repr_javascript_(self): |
| 63 | + return self._repr_mime_("application/javascript") |
| 64 | + |
| 65 | + def _repr_png_(self): |
| 66 | + return self._repr_mime_("image/png") |
| 67 | + |
| 68 | + def _repr_jpeg_(self): |
| 69 | + return self._repr_mime_("image/jpeg") |
| 70 | + |
| 71 | + def _repr_svg_(self): |
| 72 | + return self._repr_mime_("image/svg+xml") |
| 73 | + |
| 74 | + |
| 75 | +class CapturedIO(object): |
| 76 | + """Simple object for containing captured stdout/err and rich display StringIO objects |
| 77 | +
|
| 78 | + Each instance `c` has three attributes: |
| 79 | +
|
| 80 | + - ``c.stdout`` : standard output as a string |
| 81 | + - ``c.stderr`` : standard error as a string |
| 82 | + - ``c.outputs``: a list of rich display outputs |
| 83 | +
|
| 84 | + Additionally, there's a ``c.show()`` method which will print all of the |
| 85 | + above in the same order, and can be invoked simply via ``c()``. |
| 86 | + """ |
| 87 | + |
| 88 | + def __init__(self, stdout, stderr, outputs=None): |
| 89 | + self._stdout = stdout |
| 90 | + self._stderr = stderr |
| 91 | + if outputs is None: |
| 92 | + outputs = [] |
| 93 | + self._outputs = outputs |
| 94 | + |
| 95 | + def __str__(self): |
| 96 | + return self.stdout |
| 97 | + |
| 98 | + @property |
| 99 | + def stdout(self): |
| 100 | + "Captured standard output" |
| 101 | + if not self._stdout: |
| 102 | + return '' |
| 103 | + return self._stdout.getvalue() |
| 104 | + |
| 105 | + @property |
| 106 | + def stderr(self): |
| 107 | + "Captured standard error" |
| 108 | + if not self._stderr: |
| 109 | + return '' |
| 110 | + return self._stderr.getvalue() |
| 111 | + |
| 112 | + @property |
| 113 | + def outputs(self): |
| 114 | + """A list of the captured rich display outputs, if any. |
| 115 | +
|
| 116 | + If you have a CapturedIO object ``c``, these can be displayed in IPython |
| 117 | + using:: |
| 118 | +
|
| 119 | + from IPython.display import display |
| 120 | + for o in c.outputs: |
| 121 | + display(o) |
| 122 | + """ |
| 123 | + return [ RichOutput(**kargs) for kargs in self._outputs ] |
| 124 | + |
| 125 | + def show(self): |
| 126 | + """write my output to sys.stdout/err as appropriate""" |
| 127 | + sys.stdout.write(self.stdout) |
| 128 | + sys.stderr.write(self.stderr) |
| 129 | + sys.stdout.flush() |
| 130 | + sys.stderr.flush() |
| 131 | + for kargs in self._outputs: |
| 132 | + RichOutput(**kargs).display() |
| 133 | + |
| 134 | + __call__ = show |
| 135 | + |
| 136 | +class StreamLogger(object): |
| 137 | + def __init__(self, store, pass_through): |
| 138 | + self.pass_through = pass_through |
| 139 | + self.store = store |
| 140 | + |
| 141 | + def getvalue(self): |
| 142 | + return self.store.getvalue() |
| 143 | + |
| 144 | + def flush(self): |
| 145 | + self.pass_through.flush() |
| 146 | + |
| 147 | + def isatty(self): |
| 148 | + return self.pass_through.isatty() |
| 149 | + |
| 150 | + def fileno(self): |
| 151 | + return self.pass_through.fileno() |
| 152 | + |
| 153 | + def write(self, buf): |
| 154 | + self.store.write(buf) |
| 155 | + self.pass_through.write(buf) |
| 156 | + |
| 157 | + def writelines(self, lines): |
| 158 | + self.store.writelines(lines) |
| 159 | + self.pass_through.writelines(lines) |
| 160 | + |
| 161 | +class capture_output(object): |
| 162 | + """context manager for capturing stdout/err""" |
| 163 | + stdout = True |
| 164 | + stderr = True |
| 165 | + display = True |
| 166 | + |
| 167 | + def __init__(self, stdout=True, stderr=True, display=True): |
| 168 | + self.stdout = stdout |
| 169 | + self.stderr = stderr |
| 170 | + self.display = display |
| 171 | + self.shell = None |
| 172 | + |
| 173 | + def __enter__(self): |
| 174 | + from IPython.core.getipython import get_ipython |
| 175 | + from IPython.core.displaypub import CapturingDisplayPublisher |
| 176 | + from IPython.core.displayhook import CapturingDisplayHook |
| 177 | + |
| 178 | + self.sys_stdout = sys.stdout |
| 179 | + self.sys_stderr = sys.stderr |
| 180 | + |
| 181 | + if self.display: |
| 182 | + self.shell = get_ipython() |
| 183 | + if self.shell is None: |
| 184 | + self.save_display_pub = None |
| 185 | + self.display = False |
| 186 | + |
| 187 | + stdout = stderr = outputs = None |
| 188 | + if self.stdout: |
| 189 | + # stdout = sys.stdout = StringIO() |
| 190 | + stdout = StreamLogger(StringIO(), sys.stdout) |
| 191 | + sys.stdout = stdout |
| 192 | + if self.stderr: |
| 193 | + # stderr = sys.stderr = StringIO() |
| 194 | + stderr = StreamLogger(StringIO(), sys.stderr) |
| 195 | + sys.stderr = stderr |
| 196 | + if self.display: |
| 197 | + self.save_display_pub = self.shell.display_pub |
| 198 | + self.shell.display_pub = CapturingDisplayPublisher() |
| 199 | + outputs = self.shell.display_pub.outputs |
| 200 | + self.save_display_hook = sys.displayhook |
| 201 | + sys.displayhook = CapturingDisplayHook(shell=self.shell, |
| 202 | + outputs=outputs) |
| 203 | + |
| 204 | + return CapturedIO(stdout, stderr, outputs) |
| 205 | + |
| 206 | + def __exit__(self, exc_type, exc_value, traceback): |
| 207 | + sys.stdout = self.sys_stdout |
| 208 | + sys.stderr = self.sys_stderr |
| 209 | + if self.display and self.shell: |
| 210 | + self.shell.display_pub = self.save_display_pub |
| 211 | + sys.displayhook = self.save_display_hook |
| 212 | + |
| 213 | + |
| 214 | + |
| 215 | +@magics_class |
| 216 | +class MyMagics(Magics): |
| 217 | + """Magics related to code execution, debugging, profiling, etc. |
| 218 | + """ |
| 219 | + |
| 220 | + def __init__(self, shell): |
| 221 | + super(MyMagics, self).__init__(shell) |
| 222 | + # Default execution function used to actually run user code. |
| 223 | + self.default_runner = None |
| 224 | + |
| 225 | + @magic_arguments.magic_arguments() |
| 226 | + @magic_arguments.argument('output', type=str, default='', nargs='?', |
| 227 | + help="""The name of the variable in which to store output. |
| 228 | + This is a utils.io.CapturedIO object with stdout/err attributes |
| 229 | + for the text of the captured output. |
| 230 | + CapturedOutput also has a show() method for displaying the output, |
| 231 | + and __call__ as well, so you can use that to quickly display the |
| 232 | + output. |
| 233 | + If unspecified, captured output is discarded. |
| 234 | + """ |
| 235 | + ) |
| 236 | + @magic_arguments.argument('--no-stderr', action="store_true", |
| 237 | + help="""Don't capture stderr.""" |
| 238 | + ) |
| 239 | + @magic_arguments.argument('--no-stdout', action="store_true", |
| 240 | + help="""Don't capture stdout.""" |
| 241 | + ) |
| 242 | + @magic_arguments.argument('--no-display', action="store_true", |
| 243 | + help="""Don't capture IPython's rich display.""" |
| 244 | + ) |
| 245 | + @cell_magic |
| 246 | + def vsccapture(self, line, cell): |
| 247 | + """run the cell, capturing stdout, stderr, and IPython's rich display() calls.""" |
| 248 | + args = magic_arguments.parse_argstring(self.vsccapture, line) |
| 249 | + out = not args.no_stdout |
| 250 | + err = not args.no_stderr |
| 251 | + disp = not args.no_display |
| 252 | + with capture_output(out, err, disp) as io: |
| 253 | + self.shell.run_cell(cell) |
| 254 | + if DisplayHook.semicolon_at_end_of_expression(cell): |
| 255 | + if args.output in self.shell.user_ns: |
| 256 | + del self.shell.user_ns[args.output] |
| 257 | + elif args.output: |
| 258 | + self.shell.user_ns[args.output] = io |
| 259 | + |
| 260 | + |
| 261 | + |
| 262 | +# In order to actually use these magics, you must register them with a |
| 263 | +# running IPython. |
| 264 | + |
| 265 | +def load_ipython_extension(ipython): |
| 266 | + """ |
| 267 | + Any module file that define a function named `load_ipython_extension` |
| 268 | + can be loaded via `%load_ext module.path` or be configured to be |
| 269 | + autoloaded by IPython at startup time. |
| 270 | + """ |
| 271 | + # You can register the class itself without instantiating it. IPython will |
| 272 | + # call the default constructor on it. |
| 273 | + ipython.register_magics(MyMagics) |
0 commit comments