如何仅在VSCode Jupyter代码单元格中执行代码？

问题描述 #

我经常使用VSCode Jupyter代码单元格开发笔记本，同时会在tmux中通过python script.py执行脚本。目前缺少executed_as_jupyter_code_cell()函数，导致部分代码被不必要地执行，请问是否有办法检测代码是否在VSCode Jupyter代码单元格中运行？

一种变通方案是设置自定义环境变量，但或许存在更直接的解决办法？

示例代码：

#%% load

def load(x):
    ...

if executed_as_jupyter_code_cell():
    y=load(x)

#%% transform

def transform(y):
    ...

if executed_as_jupyter_code_cell():
    z=transform(y)

#%% pipe
from concurrent.futures import ProcessPoolExecutor

def pipe(x):
    y=load(x)
    z=transform(y)
    save(z)

with ProcessPoolExecutor() as exec:
    exec.map(pipe, some_long_iterable)

解决方案 #

方法1：基于IPython和VSCode环境变量检测 #

VSCode运行Jupyter代码单元格时，会自动设置VSCODE_PID、VSCODE_JUPYTER这类专属环境变量，同时代码会运行在IPython内核环境中。结合这两点可实现精准检测：

def executed_as_jupyter_code_cell():
    try:
        # 判断是否处于IPython/Jupyter环境
        from IPython import get_ipython
        ip = get_ipython()
        if ip is None:
            return False
        
        # 判断是否是VSCode的Jupyter环境
        import os
        return "VSCODE_PID" in os.environ or "VSCODE_JUPYTER" in os.environ
    except ImportError:
        # 无IPython模块，说明不是Jupyter环境
        return False

方法2：检查IPython内核标识 #

VSCode的Jupyter使用定制化IPython内核，可通过检查IPython实例的配置属性确认：

def executed_as_jupyter_code_cell():
    try:
        from IPython import get_ipython
        ip = get_ipython()
        if ip is None:
            return False
        
        # 检查内核是否来自VSCode
        return hasattr(ip, 'config') and 'VSCode' in str(ip.config.get('IPKernelApp', {}))
    except ImportError:
        return False

方法3：通用Jupyter环境检测 #

如果不需要严格区分VSCode和其他Jupyter环境（如原生Notebook），可直接检测是否处于交互式IPython环境：

def executed_as_jupyter_code_cell():
    try:
        from IPython import get_ipython
        return get_ipython() is not None
    except ImportError:
        return False

说明 #

这些方法无需手动设置自定义环境变量，直接利用VSCode和IPython运行时的内置标识，比变通方案更简洁可靠。注意：在VSCode终端中运行python script.py时，检测会返回False，符合预期。

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