Python上下文管理器

如果你有閱讀源碼的習慣,可能會看到一些優秀的代碼經常出現帶有 “with” 關鍵字的語句,它通常用在什麼場景呢?今天就來說說 with 和 上下文管理器。

對於系統資源如文件、數據庫連接、socket 而言,應用程序打開這些資源並執行完業務邏輯之後,必須做的一件事就是要關閉(斷開)該資源。

比如 Python 程序打開一個文件,往文件中寫內容,寫完之後,就要關閉該文件,否則會出現什麼情況呢?極端情況下會出現 "Too many open files" 的錯誤,因爲系統允許你打開的最大文件數量是有限的。

同樣,對於數據庫,如果連接數過多而沒有及時關閉的話,就可能會出現 "Can not connect to MySQL server Too many connections",因爲數據庫連接是一種非常昂貴的資源,不可能無限制的被創建。

來看看如何正確關閉一個文件。

普通版:

def m1():
    f = open("output.txt", "w")
    f.write("python之禪")

這樣寫有一個潛在的問題,如果在調用 write 的過程中,出現了異常進而導致後續代碼無法繼續執行,close 方法無法被正常調用,因此資源就會一直被該程序佔用而無法被釋放。那麼該如何改進代碼呢?

進階版:

def m2():
    f = open("output.txt", "w")
    try:
        f.write("python之禪")
    except IOError:
        print("oops error")
    finally:
        f.close()

改良版本的程序是對可能發生異常的代碼處進行 try 捕獲,使用 try/finally 語句,該語句表示如果在 try 代碼塊中程序出現了異常,後續代碼就不再執行,而直接跳轉到 except 代碼塊。而無論如何,finally 塊的代碼最終都會被執行。因此,只要把 close 放在 finally 代碼中,文件就一定會關閉。

高級版:


def m3():
    with open("output.txt", "w") as f:
        f.write("Python之禪")

一種更加簡潔、優雅的方式就是用 with 關鍵字。open 方法的返回值賦值給變量 f,當離開 with 代碼塊的時候,系統會自動調用 f.close() 方法, with 的作用和使用 try/finally 語句是一樣的。那麼它的實現原理是什麼?在講 with 的原理前要涉及到另外一個概念,就是上下文管理器(Context Manager)。

上下文管理器

任何實現了 __enter__() 和 __exit__() 方法的對象都可稱之爲上下文管理器,上下文管理器對象可以使用 with 關鍵字。顯然,文件(file)對象也實現了上下文管理器。

那麼文件對象是如何實現這兩個方法的呢?我們可以模擬實現一個自己的文件類,讓該類實現 __enter__() 和 __exit__() 方法。

class File():
 
    def __init__(self, filename, mode):
        self.filename = filename
        self.mode = mode
 
    def __enter__(self):
        print("entering")
        self.f = open(self.filename, self.mode)
        return self.f
 
    def __exit__(self, exc_type, exc_val, exc_tb):
        print("will exit")
        self.f.close()

(1)當執行 with 語句的時候,對象的 __enter__() 方法被觸發, 它返回的值(如果有的話)會被賦值給 as 聲明的變量。對應應該就是你將要打開的那個文件對象。

(2)然後,with 語句塊裏面的代碼開始執行。對應的就是你將要打開的那個文件對象。

(3)最後,__exit__() 方法被觸發進行清理工作。對應的輸出是【in __exit__】。

因爲 File 類實現了上下文管理器,現在就可以使用 with 語句了。


with File('out.txt', 'w') as f:
    print("writing")
    f.write('hello, python')

這樣,你就無需顯式地調用 close 方法了,由系統自動去調用,哪怕中間遇到異常,close 方法也會被調用。

補充說明:
__exit__()方法的第三個參數包含了異常類型、異常值和追溯信息(如果有的話)。 __exit__()方法能自己決定怎樣利用這個異常信息,或者忽略它並返回一個None值。

如果 __exit__() 返回 True ,那麼異常會被清空,就好像什麼都沒發生一樣, with 語句後面的程序繼續在正常執行。

上面的例子還不支持多個with嵌套使用,下面是一個可以嵌套使用with語句的例子:

from socket import socket, AF_INET, SOCK_STREAM
 
class Connection:
    def __init__(self, address, family=AF_INET, type=SOCK_STREAM):
        self.address = address
        self.family = family
        self.type = type
        self.connections = []
 
    def __enter__(self):
        sock = socket(self.family, self.type)
        sock.connect(self.address)
        self.connections.append(sock)
        return sock
 
    def __exit__(self, exc_ty, exc_val, tb):
        self.connections.pop().close()
 
 
conn = Connection(('www.python.org', 80))
with conn as s1:
    print(s1)
    with conn as s2:
        print(s2)

contextlib

Python 還提供了一個 contextmanager 的裝飾器,更進一步簡化了上下文管理器的實現方式。通過 yield 將函數分割成兩部分,yield 之前的語句在 __enter__ 方法中執行,yield 之後的語句在 __exit__ 方法中執行。緊跟在 yield 後面的值是函數的返回值。

from contextlib import contextmanager

@contextmanager
def my_open(path, mode):
    f = open(path, mode)
    yield f
    f.close()

with my_open('out.txt', 'w') as f:
    f.write("hello , the simplest context manager")

總結

Python 提供了 with 語法用於簡化資源操作的後續清除操作,是 try/finally 的替代方法,實現原理建立在上下文管理器之上。此外,Python 還提供了一個 contextmanager 裝飾器,更進一步簡化上下管理器的實現方式。

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章