Python進階——什麼是上下文管理器？

時間 2020-12-01

標籤 html python redis express app 分佈式學習優化 url spa 欄目 Python 简体版

原文原文鏈接

在 Python 開發中，咱們常常會使用到 with 語法塊，例如在讀寫文件時，保證文件描述符的正確關閉，避免資源泄露問題。html

你有沒有思考過， with 背後是如何實現的？咱們經常聽到的上下文管理器到底是什麼？python

這篇文章咱們就來學習一下 Python 上下文管理器，以及 with 的運行原理。redis

with語法塊

在講解 with 語法以前，咱們先來看一下不使用 with 的代碼如何寫？express

咱們在操做一個文件時，代碼能夠這麼寫：app

# 打開文件
f = open('file.txt')
for line in f:
    # 讀取文件內容 執行其餘操做
    # do_something...
# 關閉文件
f.close()
複製代碼

這個例子很是簡單，就是打開一個文件，而後讀取文件中的內容，最後關閉文件釋放資源。分佈式

可是，代碼這麼寫會有一個問題：在打開文件後，若是要對讀取到的內容進行其餘操做，在這操做期間發生了異常，這就會致使文件句柄沒法被釋放，進而致使資源的泄露。學習

如何解決這個問題？優化

也很簡單，咱們使用 try ... finally 來優化代碼：url

# 打開文件
f = open('file.txt')
try:
    for line in f:
        # 讀取文件內容 執行其餘操做
        # do_something...
finally:
    # 保證關閉文件
    f.close()
複製代碼

這麼寫的好處是，在讀取文件內容和操做期間，不管是否發生異常，均可以保證最後能釋放文件資源。spa

但這麼優化，代碼結構會變得很繁瑣，每次都要給代碼邏輯增長 try ... finally 才能夠，可讀性變得不好。

針對這種狀況，咱們就可使用 with 語法塊來解決這個問題：

with open('file.txt') as f:
    for line in f:
        # do_something...
複製代碼

使用 with 語法塊能夠完成以前相同的功能，並且這麼寫的好處是，代碼結構變得很是清晰，可讀性也很好。

明白了 with 的做用，那麼 with 到底是如何運行的呢？

上下文管理器

首先，咱們來看一下 with 的語法格式：

with context_expression [as target(s)]:
    with-body
複製代碼

with 語法很是簡單，咱們只須要 with 一個表達式，而後就能夠執行自定義的業務邏輯。

可是，with 後面的表達式是能夠任意寫的嗎？

答案是否認的。要想使用 with 語法塊，with 後面的的對象須要實現「上下文管理器協議」。

什麼是「上下文管理器協議」？

一個類在 Python 中，只要實現如下方法，就實現了「上下文管理器協議」：

__enter__：在進入 with 語法塊以前調用，返回值會賦值給 with 的 target
__exit__：在退出 with 語法塊時調用，通常用做異常處理

咱們來看實現了這 2 個方法的例子：

class TestContext:

    def __enter__(self):
        print('__enter__')
        return 1

    def __exit__(self, exc_type, exc_value, exc_tb):
        print('exc_type: %s' % exc_type)
        print('exc_value: %s' % exc_value)
        print('exc_tb: %s' % exc_tb)

with TestContext() as t:
    print('t: %s' % t)
    
# Output:
# __enter__
# t: 1
# exc_type: None
# exc_value: None
# exc_tb: None
複製代碼

在這個例子中，咱們定義了 TestContext 類，它分別實現了 __enter__ 和 __exit__ 方法。

這樣一來，咱們就能夠把 TestContext 當作一個「上下文管理器」來使用，也就是經過 with TestContext() as t 方式來執行。

從輸出結果咱們能夠看到，具體的執行流程以下：

__enter__ 在進入 with 語句塊以前被調用，這個方法的返回值賦給了 with 後的 t 變量
__exit__ 在執行完 with 語句塊以後被調用

若是在 with 語句塊內發生了異常，那麼 __exit__ 方法能夠拿到關於異常的詳細信息：

exc_type：異常類型
exc_value：異常對象
exc_tb：異常堆棧信息

咱們來看一個發生異常的例子，觀察 __exit__ 方法拿到的異常信息是怎樣的：

with TestContext() as t:
    # 這裏會發生異常
    a = 1 / 0 
    print('t: %s' % t)

# Output:
# __enter__
# exc_type: <type 'exceptions.ZeroDivisionError'>
# exc_value: integer division or modulo by zero
# exc_tb: <traceback object at 0x10d66dd88>
# Traceback (most recent call last):
#   File "base.py", line 16, in <module>
#     a = 1 / 0
# ZeroDivisionError: integer division or modulo by zero
複製代碼

從輸出結果咱們能夠看到，當 with 語法塊內發生異常後，__exit__ 輸出了這個異常的詳細信息，其中包括異常類型、異常對象、異常堆棧。

若是咱們須要對異常作特殊處理，就能夠在這個方法中實現自定義邏輯。

回到最開始咱們講的，使用 with 讀取文件的例子。之因此 with 可以自動關閉文件資源，就是由於內置的文件對象實現了「上下文管理器協議」，這個文件對象的 __enter__ 方法返回了文件句柄，而且在 __exit__ 中實現了文件資源的關閉，另外，當 with 語法塊內有異常發生時，會拋出異常給調用者。

僞代碼能夠這麼寫：

class File:

    def __enter__(self):
        return file_obj

    def __exit__(self, exc_type, exc_value, exc_tb):
        # with 退出時釋放文件資源
        file_obj.close()
        # 若是 with 內有異常發生 拋出異常
        if exc_type is not None:
            raise exception
複製代碼

這裏咱們小結一下，經過對 with 的學習，咱們瞭解到，with 很是適合用須要對於上下文處理的場景，例如操做文件、Socket，這些場景都須要在執行完業務邏輯後，釋放資源。

contextlib模塊

對於須要上下文管理的場景，除了本身實現 __enter__ 和 __exit__ 以外，還有更簡單的方式來作嗎？

答案是確定的。咱們可使用 Python 標準庫提供的 contextlib 模塊，來簡化咱們的代碼。

使用 contextlib 模塊，咱們能夠把上下文管理器當成一個「裝飾器」來使用。

其中，contextlib 模塊提供了 contextmanager 裝飾器和 closing 方法。

下面咱們經過例子來看一下它們是如何使用的。

contextmanager裝飾器

咱們先來看 contextmanager 裝飾器的使用：

from contextlib import contextmanager

@contextmanager
def test():
    print('before')
    yield 'hello'
    print('after')

with test() as t:
    print(t)

# Output:
# before
# hello
# after
複製代碼

在這個例子中，咱們使用 contextmanager 裝飾器和 yield配合，實現了和前面上下文管理器相同的功能，它的執行流程以下：

執行 test() 方法，先打印出 before
執行 yield 'hello'，test 方法返回，hello 返回值會賦值給 with 語句塊的 t 變量
執行 with 語句塊內的邏輯，打印出 t 的值 hello
又回到 test 方法中，執行 yield 後面的邏輯，打印出 after

這樣一來，當咱們使用這個 contextmanager 裝飾器後，就不用再寫一個類來實現上下文管理協議，只須要用一個方法裝飾對應的方法，就能夠實現相同的功能。

不過有一點須要咱們注意：在使用 contextmanager 裝飾器時，若是被裝飾的方法內發生了異常，那麼咱們須要在本身的方法中進行異常處理，不然將不會執行 yield 以後的邏輯。

@contextmanager
def test():
    print('before')
    try:
        yield 'hello'
        # 這裏發生異常 必須本身處理異常邏輯 不然不會向下執行
        a = 1 / 0 
    finally:
        print('after')

with test() as t:
    print(t)
複製代碼

closing方法

咱們再來看 contextlib 提供的 closing 方法如何使用。

closing 主要用在已經實現 close 方法的資源對象上：

from contextlib import closing

class Test():

    # 定義了 close 方法纔可使用 closing 裝飾器
    def close(self):
        print('closed')

# with 塊執行結束後 自動執行 close 方法
with closing(Test()):
    print('do something')
    
# Output:
# do something
# closed
複製代碼

從執行結果咱們能夠看到，with 語句塊執行結束後，會自動調用 Test 實例的 close 方法。

因此，對於須要自定義關閉資源的場景，咱們可使用這個方法配合 with 來完成。

contextlib的實現

學習完了 contextlib 模塊的使用，最後咱們來看一下 contextlib 模塊是到底是如何實現的？

contextlib 模塊相關的源碼以下：

class _GeneratorContextManagerBase:

    def __init__(self, func, args, kwds):
        # 接收一個生成器對象 (方法內包含 yield 的方法就是一個生成器)
        self.gen = func(*args, **kwds)
        self.func, self.args, self.kwds = func, args, kwds
        doc = getattr(func, "__doc__", None)
        if doc is None:
            doc = type(self).__doc__
        self.__doc__ = doc

class _GeneratorContextManager(_GeneratorContextManagerBase,
                               AbstractContextManager,
                               ContextDecorator):

    def __enter__(self):
        try:
            # 執行生成器 代碼會運行生成器方法的 yield 處
            return next(self.gen)
        except StopIteration:
            raise RuntimeError("generator didn't yield") from None

    def __exit__(self, type, value, traceback):
        # with 內沒有異常發生
        if type is None:
            try:
                # 繼續執行生成器
                next(self.gen)
            except StopIteration:
                return False
            else:
                raise RuntimeError("generator didn't stop")
        # with 內發生了異常
        else:
            if value is None:
                value = type()
            try:
                # 拋出異常
                self.gen.throw(type, value, traceback)
            except StopIteration as exc:
                return exc is not value
            except RuntimeError as exc:
                if exc is value:
                    return False
                if type is StopIteration and exc.__cause__ is value:
                    return False
                raise
            except:
                if sys.exc_info()[1] is value:
                    return False
                raise
            raise RuntimeError("generator didn't stop after throw()")

def contextmanager(func):
    @wraps(func)
    def helper(*args, **kwds):
        return _GeneratorContextManager(func, args, kwds)
    return helper

class closing(AbstractContextManager):
    def __init__(self, thing):
        self.thing = thing
    def __enter__(self):
        return self.thing
    def __exit__(self, *exc_info):
        self.thing.close()
複製代碼

源碼中我已經添加好了註釋，你能夠詳細看一下。

contextlib 源碼中邏輯其實比較簡單，其中 contextmanager 裝飾器實現邏輯以下：

初始化一個 _GeneratorContextManager 類，構造方法接受了一個生成器 gen
這個類實現了上下文管理器協議 __enter__ 和 __exit__
執行 with 時會進入到 __enter__ 方法，而後執行這個生成器，執行時會運行到 with 語法塊內的 yield 處
__enter__ 返回 yield 的結果
若是 with 語法塊沒有發生異常，with 執行結束後，會進入到 __exit__ 方法，再次執行生成器，這時會運行 yield 以後的代碼邏輯
若是 with 語法塊發生了異常，__exit__ 會把這個異常經過生成器，傳入到 with 語法塊內，也就是把異常拋給調用者

再來看 closing 的實現，closing 方法就是在 __exit__ 方法中調用了自定義對象的 close，這樣當 with 結束後就會執行咱們定義的 close 方法。

使用場景

學習完了上下文管理器，那麼它們具體會用在什麼場景呢？

下面我舉幾個經常使用的例子來演示下，你能夠參考一下結合本身的場景使用。

Redis分佈式鎖

from contextlib import contextmanager

@contextmanager
def lock(redis, lock_key, expire):
    try:
        locked = redis.set(lock_key, 'locked', expire)
        yield locked
    finally:
        redis.delete(lock_key)

# 業務調用 with 代碼塊執行結束後 自動釋放鎖資源
with lock(redis, 'locked', 3) as locked:
    if not locked:
        return
    # do something ...
複製代碼

在這個例子中，咱們實現了 lock 方法，用於在 Redis 上申請一個分佈式鎖，而後使用 contextmanager 裝飾器裝飾了這個方法。

以後咱們業務在調用 lock 方法時，就可使用 with 語法塊了。

with 語法塊的第一步，首先判斷是否申請到了分佈式鎖，若是申請失敗，則業務邏輯直接返回。若是申請成功，則執行具體的業務邏輯，當業務邏輯執行完成後，with 退出時會自動釋放分佈式鎖，就不須要咱們每次都手動釋放鎖了。

Redis事物和管道

from contextlib import contextmanager

@contextmanager
def pipeline(redis):
    pipe = redis.pipeline()
    try:
        yield pipe
        pipe.execute()
    except Exception as exc:
        pipe.reset()
            
# 業務調用 with 代碼塊執行結束後 自動執行 execute 方法
with pipeline(redis) as pipe:
    pipe.set('key1', 'a', 30)
    pipe.zadd('key2', 'a', 1)
    pipe.sadd('key3', 'a')
複製代碼

在這個例子中，咱們定義了 pipeline 方法，並使用裝飾器 contextmanager 讓它變成了一個上下文管理器。

以後在調用 with pipeline(redis) as pipe 時，就能夠開啓一個事物和管道，而後在 with 語法塊內向這個管道中添加命令，最後 with 退出時會自動執行 pipeline 的 execute 方法，把這些命令批量發送給 Redis 服務端。

若是在執行命令時發生了異常，則會自動調用 pipeline 的 reset 方法，放棄這個事物的執行。