巧用上下文管理器和With語句精簡程式碼
巧用上下文管理器和With語句精簡程式碼
我想你對 Python 中的 with 語句一定不陌生，尤其是在檔案的輸入輸出操作中，不過我想，大部分人可能習慣了它的使用，卻並不知道隱藏在其背後的“祕密”。
那麼，究竟 with 語句要怎麼用，與之相關的上下文管理器（context manager）是什麼，它們之間又有著怎樣的聯絡呢？這節課，我就帶你一起揭開它們的神祕面紗。
什麼是上下文管理器？
在任何一門程式語言中，檔案的輸入輸出、資料庫的連線斷開等，都是很常見的資源管理操作。但資源都是有限的，在寫程式時，我們必須保證這些資源在使用過後得到釋放，不然就容易造成資源洩露，輕者使得系統處理緩慢，重則會使系統崩潰。
光說這些概念，你可能體會不到這一點，我們可以看看下面的例子：
for x in range(10000000):
f = open('test.txt', 'w')
f.write('hello')
這裡我們一共打開了 10000000 個檔案，但是用完以後都沒有關閉它們，如果你執行該段程式碼，便會報錯：
OSError:[Errno23]Toomanyopenfilesinsystem:'test.txt'
這就是一個典型的資源洩露的例子。因為程式中同時打開了太多的檔案，佔據了太多的資源，造成系統崩潰。
為了解決這個問題，不同的程式語言都引入了不同的機制。而在 Python 中，對應的解決方式便是上下文管理器（context manager）。上下文管理器，能夠幫助你自動分配並且釋放資源，其中最典型的應用便是 with 語句
for x in range(10000000):
with open('test.txt', 'w') as f:
f.write('hello')
這樣，我們每次開啟檔案“test.txt”，並寫入‘hello’之後，這個檔案便會自動關閉，相應的資源也可以得到釋放，防止資源洩露。當然，with 語句的程式碼，也可以用下面的形式表示：
f = open('test.txt', 'w')
try:
f.write('hello')
finally:
f.close()
要注意的是，最後的 finally 尤其重要，哪怕在寫入檔案時發生錯誤異常，它也可以保證該檔案最終被關閉。不過與 with 語句相比，這樣的程式碼就顯得冗餘了，並且還容易漏寫，因此我們一般更傾向於使用 with 語句。
另外一個典型的例子，是 Python 中的 threading.lock 類。舉個例子，比如我想要獲取一個鎖，執行相應的操作，完成後再釋放，那麼程式碼就可以寫成下面這樣：
some_lock = threading.Lock()
some_lock.acquire()
try:
...
finally:
some_lock.release()
而對應的 with 語句，同樣非常簡潔：
some_lock = threading.Lock()
with somelock:
...
我們可以從這兩個例子中看到，with 語句的使用，可以簡化了程式碼，有效避免資源洩露的發生。
上下文管理器的實現
基於類的上下文管理器
瞭解了上下文管理的概念和優點後，下面我們就通過具體的例子，一起來看看上下文管理器的原理，搞清楚它的內部實現。這裡，我自定義了一個上下文管理類 FileManager，模擬 Python 的開啟、關閉檔案操作：
class FileManager:
def __init__(self, name, mode):
print('calling __init__ method')
self.name = name
self.mode = mode
self.file = None
def __enter__(self):
print('calling __enter__ method')
self.file = open(self.name, self.mode)
return self.file
def __exit__(self, exc_type, exc_val, exc_tb):
print('calling __exit__ method')
if self.file:
self.file.close()
with FileManager('test.txt', 'w') as f:
print('ready to write to file')
f.write('hello world')
## 輸出
# calling __init__ method
# calling __enter__ method
# ready to write to file
# calling __exit__ method
需要注意的是，當我們用類來建立上下文管理器時，必須保證這個類包括方法__enter__()和方法__exit__()。其中，方法__enter__()返回需要被管理的資源，方法__exit__()裡通常會存在一些釋放、清理資源的操作，比如這個例子中的關閉檔案等等。
而當我們用 with 語句，執行這個上下文管理器時：
with FileManager('test.txt', 'w') as f:
f.write('hello world')
下面這四步操作會依次發生：
-
方法
__init__()被呼叫，程式初始化物件 FileManager，使得檔名（name）是”test.txt”，檔案模式 (mode) 是’w’； -
方法
__enter__()被呼叫，檔案“test.txt”以寫入的模式被開啟，並且返回 FileManager 物件賦予變數 f； -
字串“hello world”被寫入檔案“test.txt”；
-
方法
__exit__()被呼叫，負責關閉之前開啟的檔案流。
因此，這個程式的輸出是：
calling __init__ method
calling __enter__ method
ready to write to file
calling __exit__ meth
另外，值得一提的是，方法__exit__()中的引數“exc_type, exc_val, exc_tb”，分別表示 exception_type、exception_value 和 traceback。當我們執行含有上下文管理器的 with 語句時，如果有異常丟擲，異常的資訊就會包含在這三個變數中，傳入方法__exit__()。
因此，如果你需要處理可能發生的異常，可以在__exit__()新增相應的程式碼，比如下面這樣來寫：
class Foo:
def __init__(self):
print('__init__ called')
def __enter__(self):
print('__enter__ called')
return self
def __exit__(self, exc_type, exc_value, exc_tb):
print('__exit__ called')
if exc_type:
print(f'exc_type: {exc_type}')
print(f'exc_value: {exc_value}')
print(f'exc_traceback: {exc_tb}')
print('exception handled')
return True
with Foo() as obj:
raise Exception('exception raised').with_traceback(None)
# 輸出
__init__ called
__enter__ called
__exit__ called
exc_type: <class 'Exception'>
exc_value: exception raised
exc_traceback: <traceback object at 0x1046036c8>
exception handled
這裡，我們在 with 語句中手動丟擲了異常“exception raised”，你可以看到，__exit__()方法中異常，被順利捕捉並進行了處理。不過需要注意的是，如果方法__exit__()沒有返回 True，異常仍然會被丟擲。因此，如果你確定異常已經被處理了，請在__exit__()的最後，加上“return True”這條語句。
同樣的，資料庫的連線操作，也常常用上下文管理器來表示，這裡我給出了比較簡化的程式碼：
class DBConnectionManager:
def __init__(self, hostname, port):
self.hostname = hostname
self.port = port
self.connection = None
def __enter__(self):
self.connection = DBClient(self.hostname, self.port)
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.connection.close()
with DBConnectionManager('localhost', '8080') as db_client:
與前面 FileManager 的例子類似：
-
方法
__init__()負責對資料庫進行初始化，也就是將主機名、介面（這裡是 localhost 和 8080）分別賦予變數 hostname 和 port； -
方法
__enter__()連線資料庫，並且返回物件 DBConnectionManager； -
方法
__exit__()則負責關閉資料庫的連線。
這樣一來，只要你寫完了 DBconnectionManager 這個類，那麼在程式每次連線資料庫時，我們都只需要簡單地呼叫 with 語句即可，並不需要關心資料庫的關閉、異常等等，顯然大大提高了開發的效率。
同樣的，資料庫的連線操作，也常常用上下文管理器來表示
這裡我給出了比較簡化的程式碼：
class DBConnectionManager:
def __init__(self, hostname, port):
self.hostname = hostname
self.port = port
self.connection = None
def __enter__(self):
self.connection = DBClient(self.hostname, self.port)
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.connection.close()
with DBConnectionManager('localhost', '8080') as db_client:
....
基於生成器的上下文管理器
誠然，基於類的上下文管理器，在 Python 中應用廣泛，也是我們經常看到的形式，不過 Python 中的上下文管理器並不侷限於此。除了基於類，它還可以基於生成器實現。
接下來我們來看一個例子。
比如，你可以使用裝飾器 contextlib.contextmanager，來定義自己所需的基於生成器的上下文管理器，用以支援 with 語句。
還是拿前面的類上下文管理器 FileManager 來說，我們也可以用下面形式來表示：
from contextlib import contextmanager
@contextmanager
def file_manager(name, mode):
try:
f = open(name, mode)
yield f
finally:
f.close()
with file_manager('test.txt', 'w') as f:
f.write('hello world')
這段程式碼中，函式 file_manager() 是一個生成器，當我們執行 with 語句時，便會開啟檔案，並返回檔案物件 f；當 with 語句執行完後，finally block 中的關閉檔案操作便會執行。
你可以看到，使用基於生成器的上下文管理器時，我們不再用定義__enter__()和__exit__()方法，但請務必加上裝飾器 @contextmanager，這一點新手很容易疏忽。
講完這兩種不同原理的上下文管理器後，還需要強調的是，基於類的上下文管理器和基於生成器的上下文管理器，這兩者在功能上是一致的。只不過，
-
基於類的上下文管理器更加 flexible，適用於大型的系統開發；
-
而基於生成器的上下文管理器更加方便、簡潔，適用於中小型程式。
無論你使用哪一種，請不用忘記在方法__exit__()或者是 finally block 中釋放資源，這一點尤其重要。
總結
這節課，我們先通過一個簡單的例子，瞭解了資源洩露的易發生性，和其帶來的嚴重後果，從而引入了應對方案——即上下文管理器的概念。上下文管理器，通常應用在檔案的開啟關閉和資料庫的連線關閉等場景中，可以確保用過的資源得到迅速釋放，有效提高了程式的安全性，
接著，我們通過自定義上下文管理的例項，瞭解了上下文管理工作的原理，並一起學習了基於類的上下文管理器和基於生成器的上下文管理器，這兩者的功能相同，具體用哪個，取決於你的具體使用場景。
另外，上下文管理器通常和 with 語句一起使用，大大提高了程式的簡潔度。需要注意的是，當我們用 with 語句執行上下文管理器的操作時，一旦有異常丟擲，異常的型別、值等具體資訊，都會通過引數傳入__exit__()函式中。你可以自行定義相關的操作對異常進行處理，而處理完異常後，也別忘了加上“return True”這條語句，否則仍然會丟擲異常。