如何编写可同时作为Jupyter笔记本与命令行脚本运行的asyncio兼容代码

要解决这个同时兼容Jupyter（支持顶层await）和普通Python脚本（需要asyncio.run()）的异步代码问题，核心思路是动态检测当前运行环境，然后根据环境选择对应的异步执行方式。下面是具体的实现方案：

1. 环境检测函数 #

首先写一个简单的函数来判断是否处于Jupyter/IPython环境（包括VS Code的Jupyter插件环境）：

def is_running_in_jupyter():
    try:
        from IPython import get_ipython
        # 检查是否存在IPython上下文（Jupyter会提供这个对象）
        return get_ipython() is not None
    except ImportError:
        # 普通脚本环境下会触发ImportError，返回False
        return False

2. 通用异步执行工具函数 #

基于上面的检测函数，封装一个可以在两种环境下通用的异步执行工具：

import asyncio

def run_async(coroutine):
    if is_running_in_jupyter():
        # Jupyter环境中直接返回协程，利用Jupyter的autoawait特性自动执行
        return coroutine
    else:
        # 普通脚本环境中用asyncio.run()执行协程
        return asyncio.run(coroutine)

3. 实际使用示例 #

接下来就可以在你的代码中无缝使用这个工具了，不管是Jupyter还是普通脚本都能正常运行：

# %%
import polars as pl
import numpy as np
import asyncio

# 定义你的异步函数
async def load_data_async():
    # 模拟异步操作，比如异步读取数据、调用异步API等
    await asyncio.sleep(1)
    return pl.read_parquet('my_cool.parquet')

# %%
# 通用调用方式，两种环境都兼容
df = await run_async(load_data_async())

# 后续处理逻辑
print(df.head())

为什么这个方案可行？ #

• 环境检测：通过尝试导入IPython.get_ipython()来判断环境，普通脚本中没有这个模块，会触发ImportError从而返回False；Jupyter环境下则能成功获取到IPython上下文。
• 适配执行逻辑：Jupyter的autoawait特性会自动执行顶层的协程对象，所以我们直接返回协程即可；普通脚本中则用标准的asyncio.run()来启动事件循环执行协程，避免了"asyncio.run() cannot be called from a running event loop"的错误。

关于你之前尝试方案的说明 #

• nest-asyncio：确实存在兼容性问题，尤其是和一些依赖原生事件循环的库冲突，不推荐作为通用解决方案。
• get_ipython().autoawait = False：这个设置需要在Jupyter初始化完成后生效，且部分环境可能不支持修改该配置，所以效果不稳定。
• shebang指令：只能强制用IPython运行脚本，但无法兼顾需要用普通Python解释器执行的场景，灵活性不足。

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