如何同时捕获和显示ipython(jupyter笔记本)单元输出?

时间:2020-10-28 12:10:39

标签: python jupyter-notebook ipython jupyter

我想在jupyter笔记本中同时捕获和显示单元格输出,该怎么做?

enter image description here

1 个答案:

答案 0 :(得分:2)

我从内置 %%captureimplementation 中改编了以下自定义单元格魔法:

from io import StringIO
import sys

from IPython.core import magic_arguments
from IPython.core.magic import Magics, cell_magic, magics_class
from IPython.utils.capture import CapturedIO


class Tee(StringIO):
    def __init__(self, initial_value='', newline='\n', stream=None):
        self.stream = stream
        super().__init__(initial_value, newline)
    
    def write(self, data):
        if self.stream is not None:
            self.stream.write(data)
        
        super().write(data)


class capture_and_print_output(object):
    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
        from IPython.core.displayhook import CapturingDisplayHook
        
        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
        
        stdout = stderr = outputs = None
        if self.stdout:
            stdout = sys.stdout = Tee(stream=sys.stdout)
        if self.stderr:
            stderr = sys.stderr = Tee(stream=sys.stderr)
        if self.display:
            self.save_display_pub = self.shell.display_pub
            self.shell.display_pub = CapturingDisplayPublisher()
            outputs = self.shell.display_pub.outputs
            self.save_display_hook = sys.displayhook
            sys.displayhook = CapturingDisplayHook(shell=self.shell,
                                                   outputs=outputs)
        
        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
            sys.displayhook = self.save_display_hook


@magics_class
class CustomMagics(Magics):
    @magic_arguments.magic_arguments()
    @magic_arguments.argument('output', type=str, default='', nargs='?')
    @magic_arguments.argument('--no-stderr', action='store_true')
    @magic_arguments.argument('--no-stdout', action='store_true')
    @magic_arguments.argument('--no-display', action='store_true')
    @cell_magic
    def tee(self, line, cell):
        args = magic_arguments.parse_argstring(self.tee, line)
        out = not args.no_stdout
        err = not args.no_stderr
        disp = not args.no_display
        with capture_and_print_output(out, err, disp) as io:
            self.shell.run_cell(cell)
        if args.output:
            self.shell.user_ns[args.output] = io

您需要注册此单元格魔法,通常在您的启动配置文件中或通过将其定义为您稍后加载的扩展程序,例如:

from IPython import get_ipython

get_ipython().register_magics(CustomMagics)

现在您可以按如下方式使用它:

>>> %%tee output
    print(42)
42
>>> output.show()
42

需要注意的一点:它会像原始 %%capture 一样捕获丰富的输出,但不会在原始单元格中打印它们。如果这很重要,您可能会找到一些方法来调整上述内容以包含丰富的输出,也许是通过子类化 CapturingDisplayPublisher/Hook