Python 中的关键字 with 详解

在 Python 2.5 中，with关键字被加入。它将常用的 try ... except ... finally ...模式很方便的被复用。看一个最经典的例子：

with open('file.txt') as f:
    content = f.read()

在这段代码中，无论with中的代码块在执行的过程中发生任何情况，文件最终都会被关闭。如果代码块在执行的过程中发生了一个异常，那么在这个异常被抛出前，程序会先将被打开的文件关闭。

再看另外一个例子。

在发起一个数据库事务请求的时候，经常会用类似这样的代码：

db.begin()

try:
    # do some actions
except:
    db.rollback()
    raise
finally:
    db.commit()

如果将发起事务请求的操作变成可以支持with关键字的，那么用像这样的代码就可以了：

with transaction(db):
    # do some actions

下面，详细的说明一下with的执行过程，并用两种常用的方式实现上面的代码。

with 的一般执行过程

一段基本的with表达式，其结构是这样的：

with EXPR as VAR:
    BLOCK

其中：EXPR可以是任意表达式；as VAR是可选的。其一般的执行过程是这样的：

1. 计算EXPR，并获取一个上下文管理器。
2. 上下文管理器的__exit()__方法被保存起来用于之后的调用。
3. 调用上下文管理器的__enter()__方法。
4. 如果with表达式包含as VAR，那么EXPR的返回值被赋值给VAR。
5. 执行BLOCK中的表达式。
6. 调用上下文管理器的__exit()__方法。如果BLOCK的执行过程中发生了一个异常导致程序退出，那么异常的type、value和traceback(即sys.exc_info()的返回值)将作为参数传递给__exit()__方法。否则，将传递三个None。

将这个过程用代码表示，是这样的：

mgr = (EXPR)
exit = type(mgr).__exit__ # 这里没有执行
value = type(mgr).__enter__(mgr)
exc = True

try:
    try:
        VAR = value # 如果有 as VAR
        BLOCK
    except:
        exc = False
        if not exit(mgr, *sys.exc_info()):
            raise
finally:
    if exc:
        exit(mgr, None, None, None)

这个过程有几个细节：

• 如果上下文管理器中没有__enter()__或者__exit()__中的任意一个方法，那么解释器会抛出一个AttributeError。
• 在BLOCK中发生异常后，如果__exit()__方法返回一个可被看成是True的值，那么这个异常就不会被抛出，后面的代码会继续执行。

接下来，用两种方法来实现上面来实现上面的过程的吧。

实现上下文管理器类

第一种方法是实现一个类，其含有一个实例属性db和上下文管理器所需要的方法__enter()__和__exit()__。

class transaction(object):
    def __init__(self, db):
        self.db = db
        
    def __enter__(self):
        self.db.begin()
        
    def __exit__(self, type, value, traceback):
        if type is None:
            db.commit()
        else:
            db.rollback()

了解with的执行过程后，这个实现方式是很容易理解的。下面介绍的实现方式，其原理理解起来要复杂很多。

使用生成器装饰器

在 Python 的标准库中，有一个装饰器可以通过生成器获取上下文管理器。使用生成器装饰器的实现过程如下：

from contextlib import contextmanager

@contextmanager
def transaction(db):
    db.begin()
    
    try:
        yield db
    except:
        db.rollback()
        raise
    else:
        db.commit()

第一眼上看去，这种实现方式更为简单，但是其机制更为复杂。看一下其执行过程吧：

1. Python 解释器识别到yield关键字后，def会创建一个生成器函数替代常规的函数（在类定义之外我喜欢用函数代替方法）。
2. 装饰器contextmanager被调用并返回一个帮助函数，这个帮助函数在被调用后会生成一个GeneratorContextManager实例。最终with表达式中的EXPR调用的是由contentmanager装饰器返回的帮助函数。
3. with表达式调用transaction(db)，实际上是调用帮助函数。帮助函数调用生成器函数，生成器函数创建一个生成器。
4. 帮助函数将这个生成器传递给GeneratorContextManager，并创建一个GeneratorContextManager的实例对象作为上下文管理器。
5. with表达式调用实例对象的上下文管理器的__enter()__方法。
6. __enter()__方法中会调用这个生成器的next()方法。这时候，生成器方法会执行到yield db处停止，并将db作为next()的返回值。如果有as VAR，那么它将会被赋值给VAR。
7. with中的BLOCK被执行。
8. BLOCK执行结束后，调用上下文管理器的__exit()__方法。__exit()__方法会再次调用生成器的next()方法。如果发生StopIteration异常，则pass。
9. 如果没有发生异常生成器方法将会执行db.commit()，否则会执行db.rollback()。

再次看看上述过程的代码大致实现：

def contextmanager(func):
    def helper(*args, **kwargs):
        return GeneratorContextManager(func(*args, **kwargs))
    return helper
    
class GeneratorContextManager(object):
    def __init__(self, gen):
        self.gen = gen
        
    def __enter__(self):
        try:
            return self.gen.next()
        except StopIteration:
            raise RuntimeError("generator didn't yield")
            
    def __exit__(self, type, value, traceback):
        if type is None:
            try:
                self.gen.next()
            except StopIteration:
                pass
            else:
                raise RuntimeError("generator didn't stop")
        else:
            try:
                self.gen.throw(type, value, traceback)
                raise RuntimeError("generator didn't stop after throw()")
            except StopIteration:
                return True
            except:
                if sys.exc_info()[1] is not value:
                    raise

总结

Python 的with表达式包含了很多 Python 特性，花点时间吃透with是一件非常值得的事情。

一些其他的例子

锁机制

@contextmanager
def locked(lock):
    lock.acquired()
    try:
        yield
    finally:
        lock.release()

标准输出重定向

@contextmanager
def stdout_redirect(new_stdout):
    old_stdout = sys.stdout
    sys.stdout = new_stdout
    try:
        yield
    finally:
        sys.stdout = old_stdout

with open("file.txt", "w") as f:
    with stdout_redirect(f):
        print "hello world"

引用

• 文章原地址
• The Python "with" Statement by Example
• PEP 343