Python中@contextmanager的用法

問題引出

　　在Python中with的用法一文中已經寫到了什么是上下文管理器，以及如何創建一個符合上下文協議的自定義類。那么如果我們需要將一個非自定義的類改成一個上下文管理器又該怎樣實現？

　　我們可以為需要改寫的普通類創建一個子類，在其中添加上__enter__和__exit__方法通過繼承父類的方式實現上下文管理器。但是通過contextlib模塊下提供的@contextmanager裝飾器，我們能夠更方便的將一個普通類變成上下文管理器。

 

@contextmanager

　　@contextmanager通過將一個函數變成生成器的方式來為普通類添加進入和退出時的處理代碼，從而實現了將普通類變成一個上下文管理器。

 

示例1

from contextlib import contextmanager
class File():
    def query(self):
        print('查詢文件')

@contextmanager
def open():
    print('打開文件')
    yield File()
    print('關閉文件')

with open() as f:
    f.query()
    
# 結果 打開文件 查詢文件 關閉文件

 

執行流程

　　①with語句調用open函數=>②執行open中yield之前的代碼(打開文件)=>③執行yield語句中的代碼(File())=>④執行with語句中的代碼(f.query)=>⑤執行yiled語句后的代碼(關閉文件)

 

注意事項

• 此時的with語句后是調用被修飾的函數，而不是實例化上下文管理器對象
• with open() as f 中的f是yield語句中的內容，上述示例中即所創建的File對象
• yield前、中、后的代碼都可以省略不寫

 

示例2

　　如果需要改成上下文管理器的類允許被改寫，或者使用問題引出中所提到的構建子類的方式，那么我們可以將@contextmanager所修飾的函數放在類的內部作為一個類方法來使用：

from contextlib import contextmanager
class DataBase():
    def query(self):
        print('寫入操作')
    @contextmanager
    def open(self):
        try:
            yield
            self.commit()
        except Exception as e:
            self.rollback()
            raise e

db = DataBase()
with db.open():
    db.query()

 

執行流程

　　①with語句調用open函數=>②執行try異常判斷=>③執行with語句中的寫入操作=>④執行數據庫的提交操作=>⑤如果出現異常執行數據庫回滾操作

 

 

總結

• 在思考過程中我們可以假設with語句中的內容整體移到了yield語句后，從而將with語句的執行簡化為@contextmanager所修飾的函數的執行
• 當我們需要在一段代碼的執行前后加上其他的處理代碼時都可以使用@contextmanager來完成，並不需要一定是對上下文管理器的操作