1 模塊簡介
在數年前,Python 2.5 加入了一個非常特殊的關鍵字,就是with。with語句允許開發者創建上下文管理器。什么是上下文管理器?上下文管理器就是允許你可以自動地開始和結束一些事情。例如,你可能想要打開一個文件,然后寫入一些內容,最后再關閉文件。這或許就是上下文管理器中一個最經典的示例。事實上,當你利用with語句打開一個文件時,Python替你自動創建了一個上下文管理器。
with open("test/test.txt","w") as f_obj:
f_obj.write("hello")
如果你使用的是Python 2.4,你不得不以一種老的方式來完成這個任務。
f_obj = open("test/test.txt","w")
f_obj.write("hello")
f_obj.close()
上下文管理器背后工作的機制是使用Python的方法:__enter__和__exit__。讓我們嘗試着去創建我們的上下文管理器,以此來了解上下文管理器是如何工作的。
2 模塊使用
2.1 創建一個上下文管理器類
與其繼續使用Python打開文件這個例子,不如我們創建一個上下文管理器,這個上下文管理器將會創建一個SQLite數據庫連接,當任務處理完畢,將會將其關閉。下面就是一個簡單的示例。
import sqlite3
class DataConn:
def __init__(self,db_name):
self.db_name = db_name
def __enter__(self):
self.conn = sqlite3.connect(self.db_name)
return self.conn
def __exit__(self,exc_type,exc_val,exc_tb):
self.conn.close()
if exc_val:
raise
if __name__ == "__main__":
db = "test/test.db"
with DataConn(db) as conn:
cursor = conn.cursor()
在上述代碼中,我們創建了一個類,獲取到SQLite數據庫文件的路徑。__enter__方法將會自動執行,並返回數據庫連接對象。現在我們已經獲取到數據庫連接對象,然后我們創建光標,向數據庫寫入數據或者對數據庫進行查詢。當我們退出with語句的時候,它將會調用__exit__方法用於執行和關閉這個連接。
讓我們使用其它的方法來創建上下文管理器。
2.2 利用contextlib創建一個上下文管理器
Python 2.5 不僅僅添加了with語句,它也添加了contextlib模塊。這就允許我們使用contextlib的contextmanager函數作為裝飾器,來創建一個上下文管理器。讓我們嘗試着用它來創建一個上下文管理器,用於打開和關閉文件。
from contextlib import contextmanager
@contextmanager
def file_open(path):
try:
f_obj = open(path,"w")
yield f_obj
except OSError:
print("We had an error!")
finally:
print("Closing file")
f_obj.close()
if __name__ == "__main__":
with file_open("test/test.txt") as fobj:
fobj.write("Testing context managers")
在這里,我們從contextlib模塊中引入contextmanager,然后裝飾我們所定義的file_open函數。這就允許我們使用Python的with語句來調用file_open函數。在函數中,我們打開文件,然后通過yield,將其傳遞出去,最終主調函數可以使用它。
一旦with語句結束,控制就會返回給file_open函數,它繼續執行yield語句后面的代碼。這個最終會執行finally語句--關閉文件。如果我們在打開文件時遇到了OSError錯誤,它就會被捕獲,最終finally語句依然會關閉文件句柄。
contextlib.closing(thing)
contextlib模塊提供了一些很方便的工具。第一個工具就是closing類,一旦代碼塊運行完畢,它就會將事件關閉。Python官方文檔給出了類似於以下的一個示例,
>>> from contextlib import contextmanager
>>> @contextmanager
... def closing(db):
... try:
... yield db.conn()
... finally:
... db.close()
在這段代碼中,我們創建了一個關閉函數,它被包裹在contextmanager中。這個與closing類相同。區別就是,我們可以在with語句中使用closing類本身,而非裝飾器。讓我們看如下的示例,
>>> from contextlib import closing
>>> from urllib.request import urlopen
>>> with closing(urlopen("http://www.google.com")) as webpage:
... for line in webpage:
... pass
在這個示例中,我們在closing類中打開一個url網頁。一旦我們運行完畢with語句,指向網頁的句柄就會關閉。
contextlib.suppress(*exceptions)
另一個工具就是在Python 3.4中加入的suppress類。這個上下文管理工具背后的理念就是它可以禁止任意數目的異常。假如我們想忽略FileNotFoundError異常。如果你書寫了如下的上下文管理器,那么它不會正常運行。
>>> with open("1.txt") as fobj:
... for line in fobj:
... print(line)
...
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
FileNotFoundError: [Errno 2] No such file or directory: '1.txt'
正如你所看到的,這個上下文管理器沒有處理這個異常,如果你想忽略這個錯誤,你可以按照如下方式來做,
>>> from contextlib import suppress
>>> with suppress(FileNotFoundError):
... with open("1.txt") as fobj:
... for line in fobj:
... print(line)
在這段代碼中,我們引入suppress,然后將我們要忽略的異常傳遞給它,在這個例子中,就是FileNotFoundError。如果你想運行這段代碼,你將會注意到,文件不存在時,什么事情都沒有發生,也沒有錯誤被拋出。請注意,這個上下文管理器是可重用的,2.4章節將會具體解釋。
contextlib.redirect_stdout/redirect_stderr
contextlib模塊還有一對用於重定向標准輸出和標准錯誤輸出的工具,分別在Python 3.4 和3.5 中加入。在這些工具被加入之前,如果你想對標准輸出重定向,你需要按照如下方式操作,
import sys
path = "test/test.txt"
with open(path,"w") as fobj:
sys.stdout = fobj
help(sum)
利用contextlib模塊,你可以按照如下方式操作,
from contextlib import redirect_stdout
path = "test/test.txt"
with open(path,"w") as fobj:
with redirect_stdout(fobj):
help(redirect_stdout)
在上面兩個例子中,我們均是將標准輸出重定向到一個文件。當我們調用Python的help函數,不是將信息輸出到標准輸出上,而是將信息保存到重定向的文件中。你也可以將標准輸出重定向到緩存或者從用接口如Tkinter或wxPython中獲取的文件控制類型上。
2.3 ExitStack
ExitStack是一個上下文管理器,允許你很容易地與其它上下文管理結合或者清除。這個咋聽起來讓人有些迷糊,我們來看一個Python官方文檔的例子,或許會讓我們更容易理解它。
>>> from contextlib import ExitStack
>>> filenames = ["1.txt","2.txt"]
>>> with ExitStack as stack:
... file_objects = [stack.enter_context(open(filename)) for filename in filenames]
這段代碼就是在列表中創建一系列的上下文管理器。ExitStack維護一個寄存器的棧。當我們退出with語句時,文件就會關閉,棧就會按照相反的順序調用這些上下文管理器。
Python官方文檔中關於contextlib有很多示例,你可以學習到如下的技術點:
- 從__enter__方法中捕獲異常
- 支持不定數目的上下文管理器
- 替換掉try-finally
- 其它
2.4 可重用的上下文管理器
大部分你所創建的上下文管理器僅僅只能在with語句中使用一次,示例如下:
>>> from contextlib import contextmanager
>>> @contextmanager
... def single():
... print("Yielding")
... yield
... print("Exiting context manager")
...
>>> context = single()
>>> with context:
... pass
...
Yielding
Exiting context manager
>>> with context:
... pass
...
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
File "/usr/lib/python3.4/contextlib.py", line 61, in __enter__
raise RuntimeError("generator didn't yield") from None
RuntimeError: generator didn't yield
在這段代碼中,我們創建了一個上下文管理器實例,並嘗試着在Python的with語句中運行兩次。當第二次運行時,它拋出了RuntimeError。
但是如果我們想運行上下文管理器兩次呢?我們需要使用可重用的上下文管理器。讓我們使用之前所用過的redirect_stdout這個上下文管理器作為示例,
>>> from contextlib import redirect_stdout
>>> from io import StringIO
>>> stream = StringIO()
>>> write_to_stream = redirect_stdout(stream)
>>> with write_to_stream:
... print("Write something to the stream")
... with write_to_stream:
... print("Write something else to stream")
...
>>> print(stream.getvalue())
Write something to the stream
Write something else to stream
在這段代碼中,我們創建了一個上下文管理器,它們均向StringIO(一種內存中的文件流)寫入數據。這段代碼正常運行,而沒有像之前那樣拋出RuntimeError錯誤,原因就是redirect_stdout是可重用的,允許我們可以調用兩次。當然,實際的例子將會有更多的函數調用,會更加的復雜。一定要注意,可重用的上下文管理器不一定是線程安全的。如果你需要在線程中使用它,請先仔細閱讀Python的文檔。
2.5 總結
上下文管理器很有趣,也很方便。我經常在自動測試中使用它們,例如,打開和關閉對話。現在,你應該可以使用Python內置的工具去創建你的上下文管理器。你還可以繼續閱讀Python關於contextlib的文檔,那里有很多本文沒有覆蓋到的知識。
