将上下文管理器与装饰器混合时未捕获异常

Question

我一直在为这个问题苦苦挣扎。在我正在编写的一些代码中，我需要编写一堆文件，如果需要，还可以选择创建目录树。我的想法如下：捕获异常 IOError 并且如果它的第一个参数是 ENOENT 则创建目录结构并尝试再次写入文件。

我编写了一个相对较小的重试函数，但我想将它推广到可能引发异常的“任何”代码。这一切都奏效了，直到我遇到这样的事情：

def retry(f):
    def wrapper(*args, **kwargs):
        try:
            return f(*args, **kwargs)
        except:
            print "Gotcha here!"
    return wrapper

def update(file, value):
    @contextmanager
    @retry
    def safeopen(file, mode):
        with open(file, mode) as f:
            yield f
    try:
        with safeopen(file, 'w') as f:
            f.write(value)
    except:
        print "Gotcha there!"

update( 'tests/nonexisting/dummy.txt', 'Dummy line')

我已将代码压缩到最低限度，以显示open() 引发异常时的失败情况。在这段代码中，异常仅从 update() 中的 except 块中捕获，而不是在 wrapper() 中，所以我总是得到 Gotcha there!，尽管我希望改为 Gotcha here。我试过交换@decorator 和@contextmanager 行，没办法。我已经检查并确保调用了包装器：确实如此。只是它没有捕获来自f() 的异常。

我做错了什么？

Answer 1

问题是您将 @contextmanager 装饰器与普通功能混合在一起。 @retry 装饰器是一个普通函数，但您使用它来装饰 @contextmanager 生成器 - 这不会像您期望的那样运行，因为当您调用 @contextmanager 函数时，它的函数体是并没有实际执行。相反，会返回一个 GeneratorContextManager 对象。在直接或使用with 语句调用GeneratorContextManager 的__enter__ 方法之前，不会执行函数体。

考虑这个例子：

from contextlib import contextmanager


def retry(f):
    def wrapper(*args, **kwargs):
        try:
            print("in wrapper")
            return f(*args, **kwargs)
        except:
            print "Gotcha here!"
        finally:
            print "done"
    return wrapper

@contextmanager
@retry
def safeopen(file, mode):
    print("in safe open")
    with open(file, mode) as f:
        yield f

def update(file, value):
    try:
        print("CALLING SAFE OPEN")
        with safeopen(file, 'w') as f:
            f.write(value)
    except:
        print "Gotcha there!"

update( 'tests/nonexisting/dummy.txt', 'Dummy line')

它输出：

CALLING SAFE OPEN
in wrapper
done
in safe open
Gotcha there!

如您所见，我们在进入safeopen 的主体之前退出retry 包装器，因为safeopen 是一个上下文管理器。直到GeneratorContextManager 对象实际返回，并作为with 语句的一部分进行评估，主体才被执行，但到那时为时已晚； retry 已退出。

要解决此问题，您还需要将retry 设为@contextmanager，并使用它来装饰safeopen 上下文管理器：

from contextlib import contextmanager


def retry(f):
    @contextmanager
    def wrapper(*args, **kwargs):
        try:
            print("in wrapper")
            with f(*args, **kwargs) as out:
                yield out
        except:
            print "Gotcha here!"
        finally:
            print "done"
    return wrapper

@retry
@contextmanager
def safeopen(file, mode):
    print("in safe open")
    with open(file, mode) as f:
        yield f

def update(file, value):
    print("CALLING SAFE OPEN")
    with safeopen(file, 'w') as f:
        f.write(value)

update( 'tests/nonexisting/dummy.txt', 'Dummy line')

输出：

CALLING SAFE OPEN
in wrapper
in safe open
Gotcha here!
done

编辑：

如果您颠倒装饰器的顺序，以便retry 直接装饰safeopen，您可以使retry 实现更简单一些，因为现在您正在装饰一个生成器函数，而不是一个上下文管理器：

def retry(f):
    def wrapper(*args, **kwargs):
        try:
            print("in wrapper")
            return next(f(*args, **kwargs))  # Call next on the generator object
        except:
            print "Gotcha here!"
        finally:
            print "done"
    return wrapper

@contextmanager
@retry
def safeopen(file, mode):
    print("in safe open")
    with open(file, mode) as f:
        yield f