Python语言高级特性：上下文管理器 #P005#

相信大家在编程的时候，经常会遇到这样的场景：先执行一些准备操作，然后执行自己的业务逻辑，等业务逻辑完成以后，再执行一些清理操作。例如，对文件操作，我们需要首先打开文件，然后处理文件内容，最后关闭文件。又如，当多线程程序需要访问临界资源的时候，线程首先需要获取互斥锁，当执行完成并准备退出临界区的时候，需要释放互斥锁。对于这些场景，Python中提供了上下文管理器（Context Manager），可以通过上下文管理器来控制代码块执行前的准备动作以及执行后的收尾动作。

1 上下文管理器的概念

所谓上下文管理器，就是实现了上下文管理协议的类（实现__enter__和__exit__方法）或函数（使用contextmanager装饰器）。为了使用上下文管理器，Python 2.6和Python 3.0引入了一种特殊的语句，即with语句及其可选的as子语句。with语句与上下文管理器一起作协，为Python工程师编程提供了一些便利。

上下文管理器应用于某些特殊的情景之中，典型的应用场景是打开某种资源，对资源进行处理，最后再关闭资源。可以看到，上下文管理器的作用与常见的try/finally语句的作用比较类似，都是用于确保打开的资源无论在何种情况下，都能够及时的关闭。在其他编程语言中，一般使用try/finally语句来完成。在Python中，应该优先使用上下文管理器。因为，上下文管理器可以使用更少的代码，完成同样的功能。与此同时，上下文管理器更加的灵活。例如，开发者也可以通过实现上下文管理器协议 来编写自己的对象和函数，以便应用在with语句中，简化资源的管理逻辑。接下来我们将深入讨论上下文管理器的应用场景和实现方式。

with语句形式化定义

with语句的基本格式如下：

with expression [as variable]:
    with-block

其中，expression返回了一个上下文管理器，with语句调用上下文管理器中的方法来管理资源。

2 上下文管理器的应用场景

文件

文件对象是最常见的一种上下文管理器。在所有编程语言中，编写程序处理文件，都需要在文件处理完毕以后，及时的关闭文件。在其他编程语言中，一般使用finally语句来关闭文件。在Python中，也可以使用finally语句来保证，无论在什么情况下，文件都会被关闭。如下所示：

try:
    f = open('data.txt')
    print(f.read())
finally:
    f.close()

由于文件实现了上下文管理器协议，因此，它是一个上下文管理器。with语句与上下文管理器一起管理资源。如下所示：

with open('data.txt') as f:
    print(f.read())

当我们调用open函数时，open函数会返回一个文件对象，并赋值给相应的变量。文件对象实现了上下文管理协议，因此，它是一个上下文管理器。Python利用with语句和上下文管理器来管理资源。在这个例子中，with语句执行完成以后，会保证无论是否出现异常，这个文件对象都会关闭。

锁

上下文管理器适用于所有打开资源，对资源进行处理，最后关闭资源的场景。另一个典型的应用场景是确保加锁以后，锁的释放。使用try/finally语句的代码如下：

lock = threading.Lock()
lock.acquire()
try:
    print('Lock is held')
finally:
    lock.release()

与文件管理类似，使用上下文管理器能够有效的减少代码行数，使得代码逻辑更加清晰。如下所示：

lock = Lock()
withn lock:
    print('Lock is held')

数字精度

只要对象实现了上下文管理器协议，它便是个上下文管理器，就可以应用于with语句之中。with语句与上下文管理器一起，确保某些事情（如关闭资源、释放锁）一定会发生。标准库的decimal提供了一个上下文管理器，可以临时修改数字的精度，如下所示：

import decimal

# Decimal('0.3333333333333333333333333333')
print(decimal.Decimal('1.00') / decimal.Decimal('3.00'))

with decimal.localcontext() as ctx:
    ctx.prec = 2
    #Decimal('0.33')
    x = decimal.Decimal('1.00') / decimal.Decimal('3.00')
    print(x)

3 上下文管理器协议

上下文管理器与迭代器有一些相似之处。对于迭代器，文件对象实现迭代器协议，for循环使用迭代器协议遍历文件对象。对于上下文管理器，对象实现上下文管理器协议，with语句使用上下文管理器协议访问对象。with语句的实际工作方式如下：

1. with语句中的表达式返回一个对象，该对象必须有__enter__和__exit__方法；
2. 调用对象的__enter__方法，如果as子句存在，则将__enter__函数的返回值赋值给目标对象，否则直接丢弃；
3. 执行with语句的代码块中的代码；
4. 调用对象的__exit__方法。无论是否出现异常，exit 方法依然会被调用。如果出现了异常，则将异常信息传递给 exit。如果__exit__方法返回值为False，则异常会被重新抛出。

可以看到，with语句之所以可以管理文件对象和锁对象，是因为它们实现了上下文管理协议（即__enter__方法和__exit__方法）。我们可以使用Python语言内置的dir函数查看文件对象和Lock对象的方法。如下所示，文件对象和Lock对象，都拥有__enter__方法和__exit__方法。

In [1]: f = open('/etc/passwd')

In [2]: dir(f)
Out[2]: 
['__class__',
 ......
 '__enter__',
 '__exit__',
 'xreadlines']

In [3]: import threading

In [4]: l = threading.Lock()

In [5]: dir(l)
Out[5]: 
['__class__',
 '__enter__',
 '__exit__',
  ......
 'acquire',
 'release']

4 自定义上下文管理器

我们也可以在自己的类中实现上下文管理协议，然后在with语句中使用。要实现上下文管理协议，按照Python的标准方式，定义一个类，提供__enter__和__exit__方法即可。下面给出一个非常有用的上下文管理器。

Python 3中的open函数可以在打开文件时，指定字符集编码，Python 2中的open函数则没有这个功能。因此，在Python 2中，只能使用标准库codecs来指定打开文件的字符集编码。下面的例子中定义了一个Open类，并且实现了上下文管理器协议。在Open类中，通过标准库的codecs模块，模拟Python 3的open函数。如下所示：

#!/usr/bin/python
#-*- coding: UTF-8 -*-
import codecs

class Open(object):

    def __init__(self, filename, mode, encoding="utf-8"):
        self.fp = codecs.open(filename, mode, encoding)

    def __enter__(self):
        return self.fp

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.fp.close()


data = u"上下文管理器"
with Open('data.txt', 'w') as f:
    f.write(data)

在这个例子中，我们封装了一个名为Open的类，这个类可以像Python 3中的内置的open函数一样使用，以便将中文写入到文件之中。

5 使用contextlib实现上下文管理器

在Python中，除了按照标准的协议，定义一个可以在with语句中使用的上下文管理器以外。开发者也可以使用标准库的的contextlib模块，来简化实现上下文管理器的逻辑。该模块提供了名为contextmanager的装饰器，通过该装饰器装饰的函数，就变成了一个上下文管理器，可以用在with语句之中。这种写法比标准写法更加便捷，因此，受到了更多工程师的青睐。如下所示：

#!/usr/bin/python
#-*- coding: UTF-8 -*-
import codecs
from contextlib import contextmanager

@contextmanager
def Open(filename, mode, encoding='utf-8'):
    fp = codecs.open(filename, mode, encoding)
    try:
        yield fp
    finally:
        fp.close()

data = u"上下文管理器"
with Open('data.txt', 'w') as f:
    f.write(data)

在上面这段程序中，yield表达式所在的地方，就是with语句块中的语句所要展开的地方。with语句块所抛出的任何异常，都会由yield表达式重新抛出。

6 管理多个资源

从Python 2.7和Python 3.1开始，with语句也可以使用新的逗号语法，同时使用多个上下文管理器。例如，我们要同时打开两个文件进行处理，很多工程师会写出下面这样的代码：

with open('file1', 'r') as source:
    with open('file2', 'w') as target:
        target.write(source.read())

上面这段代码，也可以简化成下面这样：

with open('file1', 'r') as source, open('file2', 'w') as target:
    target.write(source.read())

Python的with语句和上下文管理器是非常优秀的设计。充分使用上下文管理器，不但可以减少代码中的错误，也能让代码看起来更加的清晰整洁。上下文管理器并不是必须的，它的所有功能都可以使用try/finally语句实现。但是，使用上下文管理器以后代码行数更少，逻辑更加清晰。如果工程师在可以使用上下文管理器的情况下，使用了try/finally语句，将会被认为编写的代码不够Pythonic。

7 实战案例

最后，我们来看一个实战案例，即使用上下文管理器管理数据库连接。

下面的代码是一个使用MySQLdb读写MySQL数据库的完整例子，这个例子充分考虑了异常情况的处理，例如，通过上下文管理器来提交事务或回滚事务，通过try/finally语句来保证数据库连接的关闭。在这个例子中，首先删除test数据库下的student表，然后再创建一张student表，并向表中插入记录。完成以后，读取表中的记录。为了在with语句中管理数据库连接，我们需要实现上下文管理器。有两种方法实现上下文管理器，我们选择比较简洁的contextmanager。为了使用contextmanager装饰器，需要从标准库的contextlib模块中进行导入。实现上下文管理器以后，创建连接的函数如下：

from contextlib import contextmanager

import MySQLdb as db

@contextmanager
def get_conn(**kwargs):
    conn = db.connect(host=kwargs.get('host', 'localhost'), 
                        user=kwargs.get('user'),
                        passwd=kwargs.get('passwd'),
                        port=kwargs.get('port', 3306),
                        db=kwargs.get('db'))
    try:
        yield conn
    finally:
        if conn:
            conn.close()

def execute_sql(conn, sql):
    with conn as cur:
        cur.execute(sql)

def create_table(conn):
    sql_drop_table = "DROP TABLE IF EXISTS student"
    sql_create_table = """ CREATE TABLE `student` (
                           `sno` int(11) NOT NULL,
                           `sname` varchar(20) DEFAULT NULL,
                           `sage` int(11) DEFAULT NULL,
                           PRIMARY KEY (`sno`)
                           ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 """
    for sql in [sql_drop_table, sql_create_table]:
        execute_sql(conn, sql)

def insert_data(conn, sno, sname, sage):
    INSERT_FORMAT = "insert into student values({0}, '{1}', {2})"
    sql = INSERT_FORMAT.format(sno, sname, sage)
    execute_sql(conn, sql)

def main():
    with get_conn(host='127.0.0.1', user='root', passwd='root', port=3306) as conn:

        create_table(conn)
        insert_data(conn, 1, 'zhangsan', 20)

        with conn as cur:
            cur.execute("select * from student")


if __name__ == '__main__':
    main()

为了保证无论在什么情况下，数据库连接都会关闭，我们使用了try/finally语句。既然都是使用try/finally语句，那么，我们为什么还要封装一个上下文管理器呢？对于这里的get_conn函数，我们只需要定义一次，可以使用多次。我们封装成上下文管理器以后，无论get_conn函数使用多少次，都不需要再次编写try/finally相关的逻辑，只需要使用with语句来管理数据库连接即可。

8 总结

在这篇文章中，我们介绍了Python语言中一个非常重要的语言特性，即上下文管理器。上下文管理器相对于我们接下来要介绍的装饰器，并没有那么抽象，理解起来也没有什么难度。但是，正如大家看到的，使用上下文管理器以后，with语句可以实现任何try/finally语句实行的功能，而且代码更加清晰简洁。

作者介绍

赖明星，架构师、作家。现就职于腾讯，参与并主导下一代金融级数据库平台研发。有多年的 Python 开发经验和一线互联网实战经验，擅长 C、Python、Java、MySQL、Linux 等主流技术。国内知名的 Python 技术专家和 Python 技术的积极推广者，著有《Python Linux 系统管理与自动化运维》一书。