如何在Jupyter/IPython中通过函数捕获当前单元格的输出？

问题背景 #

需要实现一个capture_output函数（位于example.py模块），用于执行指定代码并捕获Jupyter Notebook/IPython会话中的富格式输出（文本、图像等），最终将输出缓存到文件。最初尝试通过run_cell_magic调用%%capture魔法命令，但执行时报错NameError: name 'output' is not defined。

错误代码示例 #

example.py中的错误实现：

from IPython import get_ipython
def capture_output(code):
    get_ipython().run_cell_magic("capture", "output", code)
    return output

调用时的报错：

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[2], line 1
----> 1 ex.capture_output("3+2")

File example.py:6, in capture_output(code)
      4 def capture_output(code):
      5     get_ipython().run_cell_magic("capture", "output", code)
----> 6     return output

NameError: name 'output' is not defined

错误原因 #

run_cell_magic("capture", "output", code)中的"output"参数是告诉%%capture魔法命令将捕获的输出存储到IPython交互式命名空间的output变量中，而非当前函数的本地作用域。直接在函数里返回output会因为变量不在本地作用域而触发NameError。

正确实现方案 #

方案1：使用IPython内置的捕获上下文管理器 #

借助IPython.utils.capture.capture_output上下文管理器，直接在函数作用域内捕获输出，无需依赖全局命名空间：

修改后的example.py：

from IPython import get_ipython
from IPython.utils.capture import capture_output as ip_capture

def capture_output(code):
    ip = get_ipython()
    # 上下文管理器捕获所有输出
    with ip_capture() as captured:
        ip.run_cell(code)
    return captured

方案2：从IPython命名空间中获取捕获对象 #

如果坚持使用%%capture魔法命令，可从IPython的用户命名空间中取出捕获的变量，用完后清理避免污染：

修改后的example.py：

from IPython import get_ipython

def capture_output(code):
    ip = get_ipython()
    # 指定临时变量名存储捕获结果
    temp_var = "_captured_temp_output"
    ip.run_cell_magic("capture", temp_var, code)
    # 从用户命名空间取出对象
    captured = ip.user_ns[temp_var]
    # 删除临时变量，避免污染命名空间
    del ip.user_ns[temp_var]
    return captured

测试验证 #

两种方案都可以正常捕获文本和图像输出：

import example as ex

# 捕获文本输出
output = ex.capture_output("3+2")
output.show()  # 输出：5

# 捕获图像输出
output = ex.capture_output("""
import matplotlib.pyplot as plt
plt.plot([1,2,3,4])
plt.show()
""")
output.show()  # 显示生成的折线图

内容的提问来源于stack exchange，提问作者Jau A