Python装饰器异常处理器无法捕获生成器函数异常，如何解决？

为复用模块中的异常处理逻辑，我实现了exception_delegator装饰器作为异常处理器：

def exception_delegator(operation):
    @wraps(operation)
    def wrapper(*args, **kwargs):
        self = args[0]
        try:
            return operation(*args, **kwargs)
        except (ClientError, ConnectTimeoutError, ReadTimeoutError) as e:
            raise MyCustomException(
                None,
                original_exception=e,
                traceback=True,
                operation_desc=self.operation_desc,
                resource_desc=self.resource_desc,
            )
        except MyCustomException:
            raise MyCustomException(
                None,
                operation_desc=self.operation_desc,
                resource_desc=self.resource_desc,
            )

    return wrapper

该处理器在普通函数和方法中正常工作，但装饰生成器方法（如下list_objects）时，无法捕获生成器内部抛出的异常，异常会直接向上传播：

@exception_delegator
    def list_objects(self):
        paginator = self.client.get_paginator('list_objects_v2')
        iterator = paginator.paginate(Bucket=self.bucket, Prefix=self.prefix)
        for response in iterator:
            yield response

我通过以下简化示例复现了该问题：

def decorator(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except ValueError:
            raise MyCustomException
    return wrapper


class MyCustomException(Exception):
    pass

def iterator():
    for i in range(4):
        if i == 3:
            raise ValueError
        print(i)
        yield i

@decorator
def gen():
    for i in iterator():
        yield i

results = list(gen())

运行输出：

0
1
2
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
~\AppData\Local\Temp\ipykernel_21476\3739375099.py in <cell line: 26>()
     24         yield i
     25 
---> 26 results = list(gen())

~\AppData\Local\Temp\ipykernel_21476\3739375099.py in gen()
     21 @decorator
     22 def gen():
---> 23     for i in iterator():
     24         yield i
     25 

~\AppData\Local\Temp\ipykernel_21476\3739375099.py in iterator()
     15     for i in range(4):
     16         if i == 3:
---> 17             raise ValueError
     18         print(i)
     19         yield i

ValueError: 

请问我忽略了什么？如何正确处理生成器的异常？

原因分析 #

当前装饰器仅在调用生成器函数、返回生成器对象的阶段做了异常捕获，但生成器的代码逻辑是在调用next()（或通过for循环、list()等迭代操作）时才会执行的。装饰器里的try-except只包裹了return func(*args,**kwargs)，这一步仅创建生成器对象，并不会执行生成器内部的代码，因此生成器迭代时抛出的异常根本没被这个try-except捕获到。

解决方案 #

需要修改装饰器，让它能接管生成器的迭代过程，在每次迭代元素时捕获异常。具体来说，装饰器的wrapper函数需要判断被装饰函数的返回值是否为生成器/迭代器，如果是，就自行迭代并在过程中捕获异常。

修改后的exception_delegator装饰器 #

from functools import wraps
from inspect import isgenerator

def exception_delegator(operation):
    @wraps(operation)
    def wrapper(*args, **kwargs):
        self = args[0]
        result = operation(*args, **kwargs)
        
        # 判断返回值是否为生成器
        if isgenerator(result):
            try:
                # 转发生成器的所有元素，同时捕获迭代时的异常
                yield from result
            except (ClientError, ConnectTimeoutError, ReadTimeoutError) as e:
                raise MyCustomException(
                    None,
                    original_exception=e,
                    traceback=True,
                    operation_desc=self.operation_desc,
                    resource_desc=self.resource_desc,
                )
            except MyCustomException as e:
                # 保留原异常上下文，补充业务描述信息
                raise MyCustomException(
                    None,
                    original_exception=e,
                    operation_desc=self.operation_desc,
                    resource_desc=self.resource_desc,
                ) from e
        else:
            # 非生成器函数，保持原有异常处理逻辑
            try:
                return result
            except (ClientError, ConnectTimeoutError, ReadTimeoutError) as e:
                raise MyCustomException(
                    None,
                    original_exception=e,
                    traceback=True,
                    operation_desc=self.operation_desc,
                    resource_desc=self.resource_desc,
                )
            except MyCustomException as e:
                raise MyCustomException(
                    None,
                    original_exception=e,
                    operation_desc=self.operation_desc,
                    resource_desc=self.resource_desc,
                ) from e
    return wrapper

简化示例的验证代码 #

把简化示例的装饰器改成如下版本，即可捕获生成器迭代时的异常：

from functools import wraps
from inspect import isgenerator

def decorator(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        result = func(*args, **kwargs)
        if isgenerator(result):
            try:
                yield from result
            except ValueError:
                raise MyCustomException
        else:
            try:
                return result
            except ValueError:
                raise MyCustomException
    return wrapper


class MyCustomException(Exception):
    pass

def iterator():
    for i in range(4):
        if i == 3:
            raise ValueError
        print(i)
        yield i

@decorator
def gen():
    for i in iterator():
        yield i

try:
    results = list(gen())
except MyCustomException as e:
    print(f"捕获到自定义异常: {e}")

运行后输出：

0
1
2
捕获到自定义异常: 

关键说明 #

• 使用yield from可以简化生成器的迭代转发逻辑，同时自动处理生成器的所有迭代细节
• 通过isgenerator判断返回值类型，保证装饰器同时兼容普通函数和生成器函数
• 对于已有的MyCustomException，使用raise ... from e保留原异常的上下文，方便后续调试排查问题

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