with 语句和上下文管理器

  1. # create/aquire some resource
  2. ...
  3. try:
  4.     # do something with the resource
  5.     ...
  6. finally:
  7.     # destroy/release the resource
  8.     ...

处理文件，线程，数据库，网络编程等等资源的时候，我们经常需要使用上面这样的代码形式，以确保资源的正常使用和释放。

好在Python 提供了 with 语句帮我们自动进行这样的处理，例如之前在打开文件时我们使用：

In [1]:

  1. with open('my_file', 'w') as fp:
  2.     # do stuff with fp
  3.     data = fp.write("Hello world")

这等效于下面的代码，但是要更简便：

In [2]:

  1. fp = open('my_file', 'w')
  2. try:
  3.     # do stuff with f
  4.     data = fp.write("Hello world")
  5. finally:
  6.     fp.close()

上下文管理器

其基本用法如下：

  1. with <expression>:
  2.     <block>

<expression> 执行的结果应当返回一个实现了上下文管理器的对象，即实现这样两个方法，enter 和 exit：

In [3]:

  1. print fp.__enter__
  2. print fp.__exit__
  1. <built-in method __enter__ of file object at 0x0000000003C1D540>
  2. <built-in method __exit__ of file object at 0x0000000003C1D540>

enter 方法在 <block> 执行前执行，而 exit 在 <block> 执行结束后执行：

比如可以这样定义一个简单的上下文管理器：

In [4]:

  1. class ContextManager(object):
  2.  
  3.     def __enter__(self):
  4.         print "Entering"
  5.  
  6.     def __exit__(self, exc_type, exc_value, traceback):
  7.         print "Exiting"

使用 with 语句执行：

In [5]:

  1. with ContextManager():
  2.     print "  Inside the with statement"
  1. Entering
  2.   Inside the with statement
  3. Exiting

即使 <block> 中执行的内容出错，exit 也会被执行：

In [6]:

  1. with ContextManager():
  2.     print 1/0
  1. Entering
  2. Exiting
  1. ---------------------------------------------------------------------------
  2. ZeroDivisionError                         Traceback (most recent call last)
  3. <ipython-input-6-b509c97cb388> in <module>()
  4.       1 with ContextManager():
  5. ----> 2     print 1/0
  6.  
  7. ZeroDivisionError: integer division or modulo by zero

enter 的返回值

如果在 enter 方法下添加了返回值，那么我们可以使用 as 把这个返回值传给某个参数：

In [7]:

  1. class ContextManager(object):
  2.  
  3.     def __enter__(self):
  4.         print "Entering"
  5.         return "my value"
  6.  
  7.     def __exit__(self, exc_type, exc_value, traceback):
  8.         print "Exiting"

将 enter 返回的值传给 value 变量：

In [8]:

  1. with ContextManager() as value:
  2.     print value
  1. Entering
  2. my value
  3. Exiting

一个通常的做法是将 enter 的返回值设为这个上下文管理器对象本身，文件对象就是这样做的：

In [9]:

  1. fp = open('my_file', 'r')
  2. print fp.__enter__()
  3. fp.close()
  1. <open file 'my_file', mode 'r' at 0x0000000003B63030>

In [10]:

  1. import os
  2. os.remove('my_file')

实现方法非常简单：

In [11]:

  1. class ContextManager(object):
  2.  
  3.     def __enter__(self):
  4.         print "Entering"
  5.         return self
  6.  
  7.     def __exit__(self, exc_type, exc_value, traceback):
  8.         print "Exiting"

In [12]:

  1. with ContextManager() as value:
  2.     print value
  1. Entering
  2. <__main__.ContextManager object at 0x0000000003D48828>
  3. Exiting

错误处理

上下文管理器对象将错误处理交给 exit 进行，可以将错误类型，错误值和 traceback 等内容作为参数传递给 exit 函数：

In [13]:

  1. class ContextManager(object):
  2.  
  3.     def __enter__(self):
  4.         print "Entering"
  5.  
  6.     def __exit__(self, exc_type, exc_value, traceback):
  7.         print "Exiting"
  8.         if exc_type is not None:
  9.             print "  Exception:", exc_value

如果没有错误，这些值都将是 None, 当有错误发生的时候：

In [14]:

  1. with ContextManager():
  2.     print 1/0
  1. Entering
  2. Exiting
  3.   Exception: integer division or modulo by zero
  1. ---------------------------------------------------------------------------
  2. ZeroDivisionError                         Traceback (most recent call last)
  3. <ipython-input-14-b509c97cb388> in <module>()
  4.       1 with ContextManager():
  5. ----> 2     print 1/0
  6.  
  7. ZeroDivisionError: integer division or modulo by zero

在这个例子中，我们只是简单的显示了错误的值，并没有对错误进行处理，所以错误被向上抛出了，如果不想让错误抛出，只需要将 exit 的返回值设为 True：

In [15]:

  1. class ContextManager(object):
  2.  
  3.     def __enter__(self):
  4.         print "Entering"
  5.  
  6.     def __exit__(self, exc_type, exc_value, traceback):
  7.         print "Exiting"
  8.         if exc_type is not None:
  9.             print " Exception suppresed:", exc_value
  10.             return True

In [16]:

  1. with ContextManager():
  2.     print 1/0
  1. Entering
  2. Exiting
  3.  Exception suppresed: integer division or modulo by zero

在这种情况下，错误就不会被向上抛出。

数据库的例子

对于数据库的 transaction 来说，如果没有错误，我们就将其 commit 进行保存，如果有错误，那么我们将其回滚到上一次成功的状态。

In [17]:

  1. class Transaction(object):
  2.  
  3.     def __init__(self, connection):
  4.         self.connection = connection
  5.  
  6.     def __enter__(self):
  7.         return self.connection.cursor()
  8.  
  9.     def __exit__(self, exc_type, exc_value, traceback):
  10.         if exc_value is None:
  11.             # transaction was OK, so commit
  12.             self.connection.commit()
  13.         else:
  14.             # transaction had a problem, so rollback
  15.             self.connection.rollback()

建立一个数据库，保存一个地址表：

In [18]:

  1. import sqlite3 as db
  2. connection = db.connect(":memory:")
  3.  
  4. with Transaction(connection) as cursor:
  5.     cursor.execute("""CREATE TABLE IF NOT EXISTS addresses (
  6.         address_id INTEGER PRIMARY KEY,
  7.         street_address TEXT,
  8.         city TEXT,
  9.         state TEXT,
  10.         country TEXT,
  11.         postal_code TEXT
  12.     )""")

插入数据：

In [19]:

  1. with Transaction(connection) as cursor:
  2.     cursor.executemany("""INSERT OR REPLACE INTO addresses VALUES (?, ?, ?, ?, ?, ?)""", [
  3.         (0, '515 Congress Ave', 'Austin', 'Texas', 'USA', '78701'),
  4.         (1, '245 Park Avenue', 'New York', 'New York', 'USA', '10167'),
  5.         (2, '21 J.J. Thompson Ave.', 'Cambridge', None, 'UK', 'CB3 0FA'),
  6.         (3, 'Supreme Business Park', 'Hiranandani Gardens, Powai, Mumbai', 'Maharashtra', 'India', '400076'),
  7.     ])

假设插入数据之后出现了问题：

In [20]:

  1. with Transaction(connection) as cursor:
  2.     cursor.execute("""INSERT OR REPLACE INTO addresses VALUES (?, ?, ?, ?, ?, ?)""",
  3.         (4, '2100 Pennsylvania Ave', 'Washington', 'DC', 'USA', '78701'),
  4.     )
  5.     raise Exception("out of addresses")
  1. ---------------------------------------------------------------------------
  2. Exception                                 Traceback (most recent call last)
  3. <ipython-input-20-ed8abdd56558> in <module>()
  4.       3         (4, '2100 Pennsylvania Ave', 'Washington', 'DC', 'USA', '78701'),
  5.       4     )
  6. ----> 5     raise Exception("out of addresses")
  7.  
  8. Exception: out of addresses

那么最新的一次插入将不会被保存，而是返回上一次 commit 成功的状态：

In [21]:

  1. cursor.execute("SELECT * FROM addresses")
  2. for row in cursor:
  3.     print row
  1. (0, u'515 Congress Ave', u'Austin', u'Texas', u'USA', u'78701')
  2. (1, u'245 Park Avenue', u'New York', u'New York', u'USA', u'10167')
  3. (2, u'21 J.J. Thompson Ave.', u'Cambridge', None, u'UK', u'CB3 0FA')
  4. (3, u'Supreme Business Park', u'Hiranandani Gardens, Powai, Mumbai', u'Maharashtra', u'India', u'400076')

contextlib 模块

很多的上下文管理器有很多相似的地方，为了防止写入很多重复的模式，可以使用 contextlib 模块来进行处理。

最简单的处理方式是使用 closing 函数确保对象的 close() 方法始终被调用：

In [23]:

  1. from contextlib import closing
  2. import urllib
  3.  
  4. with closing(urllib.urlopen('http://www.baidu.com')) as url:
  5.     html = url.read()
  6.  
  7. print html[:100]
  1. <!DOCTYPE html><!--STATUS OK--><html><head><meta http-equiv="content-type" content="text/html;charse

另一个有用的方法是使用修饰符 @contextlib：

In [24]:

  1. from contextlib import contextmanager
  2.  
  3. @contextmanager
  4. def my_contextmanager():
  5.     print "Enter"
  6.     yield
  7.     print "Exit"
  8.  
  9. with my_contextmanager():
  10.     print "  Inside the with statement"
  1. Enter
  2.   Inside the with statement
  3. Exit

yield 之前的部分可以看成是 enter 的部分，yield 的值可以看成是 enter 返回的值，yield 之后的部分可以看成是 exit 的部分。

使用 yield 的值：

In [25]:

  1. @contextmanager
  2. def my_contextmanager():
  3.     print "Enter"
  4.     yield "my value"
  5.     print "Exit"
  6.  
  7. with my_contextmanager() as value:
  8.     print value
  1. Enter
  2. my value
  3. Exit

错误处理可以用 try 块来完成：

In [26]:

  1. @contextmanager
  2. def my_contextmanager():
  3.     print "Enter"
  4.     try:
  5.         yield
  6.     except Exception as exc:
  7.         print "   Error:", exc
  8.     finally:
  9.         print "Exit"

In [27]:

  1. with my_contextmanager():
  2.     print 1/0
  1. Enter
  2.    Error: integer division or modulo by zero
  3. Exit

对于之前的数据库 transaction 我们可以这样定义：

In [28]:

  1. @contextmanager
  2. def transaction(connection):
  3.     cursor = connection.cursor()
  4.     try:
  5.         yield cursor
  6.     except:
  7.         connection.rollback()
  8.         raise
  9.     else:
  10.         connection.commit()

原文: https://nbviewer.jupyter.org/github/lijin-THU/notes-python/blob/master/05-advanced-python/05.11-context-managers-and-the-with-statement.ipynb