如何实现PEP-342风格协程，使两端挂起切换形似函数调用？

背景与需求 #

PEP-342定义的生成器协程通过yield和.send()实现上下文切换与消息传递，示例代码如下：

def coroutine(lst):
    x = yield
    lst.append(f"{x=}")
    y = yield
    lst.append(f"{y=}")
    yield

def func():
    lst = []
    coro = coroutine(lst)
    next(coro)  # 预激生成器
    coro.send(1)
    coro.send(2)
    return lst

print(func())  # 输出 ['x=1', 'y=2']

需求是将上述机制中的yield替换为函数调用get_thing()，同时隐藏预激操作（next()），让调用端直接通过.send()传递值，目标代码结构如下：

@magical_coroutine_decorator
def coroutine(lst):
    x = get_thing()
    lst.append(f"{x=}")
    y = get_thing()
    lst.append(f"{y=}")

def func():
    lst = []
    coro = coroutine(lst)
    coro.send(1)
    coro.send(2)
    return lst

手动实现方案 #

1. 核心思路 #

• 用上下文变量让get_thing()能访问当前协程的生成器实例
• 装饰器负责将原函数包装为生成器、预激生成器，并封装.send()接口
• get_thing()本质是yield的语法糖，触发生成器挂起并接收.send()传递的值

2. 代码实现 #

import contextvars

# 上下文变量：存储当前活跃的协程生成器
_current_coro = contextvars.ContextVar("_current_coro")

def get_thing():
    # 触发生成器挂起，获取调用端send的值
    return _current_coro.get().send(None)

def magical_coroutine_decorator(func):
    def wrapper(*args, **kwargs):
        # 定义内部生成器，执行原函数并绑定上下文
        def coro_gen():
            _current_coro.set(gen)
            func(*args, **kwargs)
        
        gen = coro_gen()
        next(gen)  # 预激生成器，执行到第一个get_thing()
        
        # 封装协程对象，提供send方法
        class CoroWrapper:
            def __init__(self, gen):
                self.gen = gen
            
            def send(self, value):
                try:
                    self.gen.send(value)
                except StopIteration:
                    # 协程执行完毕时忽略终止异常
                    pass
        
        return CoroWrapper(gen)
    return wrapper

3. 测试验证 #

运行目标代码，输出与PEP-342示例一致：

@magical_coroutine_decorator
def coroutine(lst):
    x = get_thing()
    lst.append(f"{x=}")
    y = get_thing()
    lst.append(f"{y=}")

def func():
    lst = []
    coro = coroutine(lst)
    coro.send(1)
    coro.send(2)
    return lst

print(func())  # 输出 ['x=1', 'y=2']

现成库方案 #

第三方库greenlet专门用于轻量级上下文切换，可更简洁实现该需求：

1. 依赖安装 #

pip install greenlet

2. 代码实现 #

from greenlet import greenlet

def get_thing():
    # 切换回调用者上下文，获取传递的值
    return greenlet.getcurrent().parent.switch()

def magical_coroutine_decorator(func):
    def wrapper(*args, **kwargs):
        # 创建协程greenlet
        coro = greenlet(lambda: func(*args, **kwargs))
        coro.switch()  # 预激协程
        
        # 封装send接口
        class CoroWrapper:
            def __init__(self, coro):
                self.coro = coro
            
            def send(self, value):
                self.coro.switch(value)
        
        return CoroWrapper(coro)
    return wrapper

3. 测试效果 #

与手动实现完全一致，greenlet的switch()方法天然支持上下文间的消息传递，无需手动处理生成器的StopIteration异常。

内容的提问来源于stack exchange，提问作者J E K