答案 0 :(得分:2)
我从内置 %%capture
的 implementation 中改编了以下自定义单元格魔法:
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
。