吃透Python上下文管理器

什麼是上下文管理器？

咱們常見的with open操做文件，就是一個上下文管理器。如：

with open(file, 'rb') as f:
    text = f.read()

那上下文管理器具體的定義是什麼呢？

上下文管理器：是指在一段代碼執行以前執行一段代碼，用於一些預處理工做；執行以後再執行一段代碼，用於一些清理工做。

好比剛提到的文件操做，打開文件進行讀寫，讀寫完以後須要將文件關閉。很明顯用到了上下文管理器。主要依靠__enter__、__exit__這兩個」魔術方法」實現。

__enter__(self)

Defines what the context manager should do at the beginning of the block created by the with statement. Note that the return value of enter is bound to the target of the with statement, or the name after the as.

__exit__(self, exception_type, exception_value, traceback)

Defines what the context manager should do after its block has been executed (or terminates). It can be used to handle exceptions, perform cleanup, or do something always done immediately after the action in the block. If the block executes successfully, exception_type, exceptionvalue, and traceback will be None. Otherwise, you can choose to handle the exception or let the user handle it; if you want to handle it, make sure _exit__ returns True after all is said and done. If you don’t want the exception to be handled by the context manager, just let it happen.

當咱們須要建立一個上下文管理器類型的時候，就須要實現__enter__和__exit__方法，這對方法稱爲上下文管理協議（Context Manager Protocol），定義了一種運行時上下文環境。

基本語法

with EXPR as VAR:
    BLOCK

這裏就是一個標準的上下文管理器的使用邏輯，其中的運行邏輯：

（1）執行EXPR語句，獲取上下文管理器（Context Manager）

（2）調用上下文管理器中的__enter__方法，該方法執行一些預處理工做。

（3）這裏的as VAR能夠省略，若是不省略，則將__enter__方法的返回值賦值給VAR。

（4）執行代碼塊BLOCK，這裏的VAR能夠當作普通變量使用。

（5）最後調用上下文管理器中的的__exit__方法。

（6）__exit__方法有三個參數：exc_type, exc_val, exc_tb。若是代碼塊BLOCK發生異常並退出，那麼分別對應異常的type、value 和 traceback。不然三個參數全爲None。

（7）__exit__方法的返回值能夠爲True或者False。若是爲True，那麼表示異常被忽視，至關於進行了try-except操做；若是爲False，則該異常會被從新raise。

如何本身實現上下文管理器？

簡單來講，若是一個類中，實現了__enter__和__exit__方法，那麼這個類就是上下文管理器。

class Contextor():
    def __enter__(self):    
        print('程序的預處理開始啦！')
        return self     # 做爲as說明符指定的變量的值

    def __exit__(self, exc_type, exc_val, exc_tb):
        print('正在進行收尾！')

    def func(self):
        print('程序進行中....')


with Contextor() as var:
    var.func()
   
# 輸出
程序的預處理開始啦！
程序進行中....
正在進行收尾！

從這個示例能夠很明顯的看出，在編寫代碼時，能夠將資源的鏈接或者獲取放在__enter__中，而將資源的關閉寫在__exit__ 中。

爲何要使用上下文管理器？

在我看來，這和 Python 崇尚的優雅風格有關。

1. 能夠以一種更加優雅的方式，操做（建立/獲取/釋放）資源，如文件操做、數據庫鏈接；
2. 能夠以一種更加優雅的方式，處理異常；

第二種，會被大多數人所忽略。這裏着重介紹下。

在處理異常時，一般都是使用try...except...來進行異常處理的。這就可能會出如今程序的主邏輯中有大量的異常處理代碼，這會大大影響程序的可讀性。

好的作法能夠經過with將異常處理隱藏起來。

咱們以1/0舉例（1/0必然會拋出錯誤）

class Resource():
    def __enter__(self):
        print('===connect to resource===')
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        print('===close resource connection===')
        return True

    def operate(self):
        1/0

with Resource() as res:
    res.operate()
    
# 輸出結果
===connect to resource===
===close resource connection===

運行發現，並無出異常。

這就是上下文管理器的強大之處，異常能夠在__exit__進行捕獲，並本身決定如何處理。在__exit__ 裏返回 True（沒有return 就默認爲 return False），就至關於告訴 Python解釋器，這個異常咱們已經捕獲了，不須要再往外拋了。

在 寫__exit__ 函數時，須要注意的事，它必需要有這三個參數：

• exc_type：異常類型
• exc_val：異常值
• exc_tb：異常的錯誤棧信息

當主邏輯代碼沒有報異常時，這三個參數將都爲None。

如何處理自行處理異常

咱們以上面的代碼爲例，在__exit__加入判斷異常的邏輯，若是發生異常，則打印異常信息。

class Resource():
    def __enter__(self):
        print('===connect to resource===')
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_type:
            print(f'出現異常：{exc_type}：{exc_val}')
        print('===close resource connection===')
        return True

    def operate(self):
        1/0

with Resource() as res:
    res.operate()
    
# 輸出
===connect to resource===
出現異常：<class 'ZeroDivisionError'>：division by zero
===close resource connection===

如何更好的使用上下文管理器？

在Python中有一個專門用於實現上下文管理的標準庫contextlib。

有了 contextlib 建立上下文管理的最好方式就是使用 contextmanager 裝飾器，經過 contextmanager 裝飾一個生成器函數，yield 語句前面的部分被認爲是__enter__() 方法的代碼，後面的部分被認爲是 __exit__()方法的代碼。

咱們以打開文件爲例：

import contextlib

@contextlib.contextmanager
def open_func(file_name):
    # __enter__方法
    print('open file:', file_name, 'in __enter__')
    file_handler = open(file_name, 'r')
    
    # 【重點】：yield
    yield file_handler

    # __exit__方法
    print('close file:', file_name, 'in __exit__')
    file_handler.close()
    return

with open_func('/Users/MING/mytest.txt') as file_in:
    for line in file_in:
        print(line)

若是要處理異常，將上面代碼改寫成下面的樣子。

import contextlib

@contextlib.contextmanager
def open_func(file_name):
    # __enter__方法
    print('open file:', file_name, 'in __enter__')
    file_handler = open(file_name, 'r')

    try:
        yield file_handler
    except Exception as exc:
        # deal with exception
        print('the exception was thrown')
    finally:
        print('close file:', file_name, 'in __exit__')
        file_handler.close()

        return

with open_func('test.txt') as file_in:
    for line in file_in:
        1/0
        print(line)

參考：

https://juejin.im/post/5c87b165f265da2dac4589cc
https://www.cnblogs.com/linxiyue/p/3855751.html
https://runnerliu.github.io/2018/01/02/pythoncontextmanager/

掃描下面二維碼，關注公衆號， 每週按期與您分享原創的、有深度的Python知識點。