Python巧用上下文管理器和with語句精簡代碼

我們在Python中對於with的語句應該是不陌生的,特別是在文件的輸入輸出操作中,那在具體的使用過程中,是有什麼引伸的含義呢?與之密切相關的上下文管理器(context manager)又是什麼呢?

什麼是上下文管理器

在任何一種編程語言裏,文件的輸入輸出、數據庫的建立連接和斷開等操作,都是很常見的資源管理操作。但是資源是有限的,在寫程序的時候,我們必須保證這些資源在使用後得到釋放,不然就容易造成資源泄漏,輕者系統處理緩慢,重則系統崩潰。

我們看一個例子:

for i in range(100000000):
    f = open('test.txt','w')
    f.write('hello')

我們在循環裏打開了100000000個文件,但是在使用完畢後沒有進行關閉操作,一運行代碼,就報錯了。

這就是一個典型的資源泄漏的案例,因爲程序中同時打開了太多的文件,佔用了太多的資源,造成崩潰。

爲了解決這個問題,不同的編程語言都引入了不同的機制,在Python中,對應的解決方法就是上下文管理器(context manager)。上下文管理器能夠自動分配資源並釋放資源,其中最典型的應用就是with語句,所以上面的代碼應該用這種方式來寫

'''
遇到問題沒人解答?小編創建了一個Python學習交流QQ羣:579817333 
尋找有志同道合的小夥伴,互幫互助,羣裏還有不錯的視頻學習教程和PDF電子書!
'''
for x in range(100000000):
    with open('test.txt','w') as f:
        f.write('hello world')

這樣,我們每次打開文件“test.txt"並寫入字符以後,這個文件就會自動關閉,相應的資源也就得以釋放,可以防止資源泄漏。當然,with的語句也可以用下面的方式來表示

f = open('test.txt','w')
try:
    f.write('hello world')
finally:
    f.close()

這裏一定要注意finally的程序段,哪怕在寫入的時候發生了異常,他也可以保證文件最終被關閉。不過於with想比較,就顯得比較冗餘了,並且還容易忽略finally,所以我們平時更傾向於使用with語句。

另外一種很典型的例子,就是Python中的線程鎖(threrading.lock類),比如我們想要獲得一個鎖,執行相應的操作以後再將其釋放,那麼代碼就應該是這樣的

import threading
some_lock = threading.Lock()
some_lock.acquire()
try:
    pass
finally:
    some_lock.release()

而與其對應的with語句就非常簡潔了

import threading
some_lock = threading.Lock()
with some_lock:
    pass

從上面兩個例子可以發現,使用with語句,可以大大的簡化代碼結構,有效的避免資源泄漏的發生。

上下文管理器的實現

基於類的上下文管理器

了接了上下文管理的概念和優點以後,我們就通過下面的例子,看看上下文管理器的原理,高清他的內部實現。我們在這裏定義一個上下文管理類FileManager,來模擬Python的打開、關閉文件的操作

'''
遇到問題沒人解答?小編創建了一個Python學習交流QQ羣:579817333 
尋找有志同道合的小夥伴,互幫互助,羣裏還有不錯的視頻學習教程和PDF電子書!
'''
class FileManager():
    def __init__(self,name,mode):
        print('call __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('call __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')

##########輸出##########
call __init__ method
calling __enter__ method
ready to write to file
call __exit__ method

特別注意:當我們用類來創建上下文管理器的時候,必須保證這個類包括下面兩個方法:

__enter__()
__exit__()

並且enter方法還要返回需要被管理的資源,方法exit裏通常會存在一些釋放、清理資源的操作,比如上面這段代碼裏的關閉文件等。

而當我們用with語句來執行上面這個上下文管理器的時候,會發生下面四個步驟:

1.構造方法__init__()會被調用,程序初始化對象FileManager,使得文件名和操作方式被傳入。

2.方法__enter__()被調用,文件被以寫入的模式打開,並且返回FileManager對戲那個賦值給變量f

3.字符串被寫入文件

4.方法__exit__()被調用,關閉之前打開的文件流。

所以就有了上面列出的輸出結果。

另外我們可以看到exit函數裏傳遞了幾個參數——exc_type.exc_val,exc_tb,分別表示exception_type,exception_value和traceback。當我們執行含有上下文管理器的with語句的時候,如果有異常拋出,異常的信息就會被包含在上面三個參數中,傳給__exit__()函數。

因此,如果我們需要處理一些異常,可以在__exit__()函數中添加相應的代碼

'''
遇到問題沒人解答?小編創建了一個Python學習交流QQ羣:579817333 
尋找有志同道合的小夥伴,互幫互助,羣裏還有不錯的視頻學習教程和PDF電子書!
'''
def __exit__(self,exc_type,exc_val,exc_tb):
        print('call __exit__ method')
        if exc_type:
            print(f'exc_type:{exc_type}')
            print(f'exc_value:{exc_val}')
            print(f'exc_traceback:{exc_tb}')
            print('exception handled')
            return True
        if self.file:
            self.file.close()

with FileManager('test.txt','w') as f:
    raise Exception('exception raised').with_traceback(None)

在修改了exit()方法以後我們在with語句中用raise手動拋出異常,我們可以看到代碼有下面的輸出

call __init__ method
calling __enter__ method
call __exit__ method
exc_type:<class 'Exception'>
exc_value:exception raised
exc_traceback:<traceback object at 0x000001AA82C8FF48>
exception handled

要注意的是,如果exit函數如果沒有返回True,呢麼異常仍然會被拋出。如果我們確定了異常已經被處理,那麼在exit最後要加上True的返回值。

同樣的,在數據庫的連接等操作上,也常常使用上下文管理器來表示,下面是個簡化的代碼

'''
遇到問題沒人解答?小編創建了一個Python學習交流QQ羣:579817333 
尋找有志同道合的小夥伴,互幫互助,羣裏還有不錯的視頻學習教程和PDF電子書!
'''
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:
    pass

代碼的具體含義和上面的例子類似,就不再詳細說明了。只要我們寫完了DBConnectionManager這個類以後,在每次建立數據庫連接時,只要簡單的利用with語句就可以了,並不要關係數據庫的關閉、異常等,大大的提高了開發效率。

基於生成器的上下文管理器

上面那種基於類的上下文管理器在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代碼段中的關閉操作就會執行。

可以看到,使用基於生成器的上下文管理器的時候,我們不用再定義__enter__()和__exit__()兩個函數,但是必須加上裝飾器,這一點非常容易漏掉。

講完這兩種上下文管理器以後,我們要強調一點:不論是基於類的還是生成器的上下文管理器,兩者在功能上是一樣的,只不過有下面兩點:

1.基於類的上下文管理器更加靈活,適用於大型的系統開發

2.基於生成器的上下文管理器更加方便、簡潔,適用於中小型程序。

但是無論使用哪一種,我們一定要記得在exit函數或finally裏寫好釋放資源的代碼,這一點尤爲重要。

總結

在這一章的開頭我們通過一個簡單的例子瞭解了資源泄漏的易發生的特性和其帶來的後果,從而引入了上下文管理器這個概念:

上下文管理器通常用在文件的IO操作和數據庫的連接關閉等場景中,可以確保用過的資源得到迅速釋放,有效提高了程序的安全性。

接着,我們通過自定義上下文管理器的示例,大致瞭解了上下文管理器工作的原理,並介紹基於類的上下文管理器和基於生成器的上下文管理器:兩者功能相同,具體使用哪個要根據場景來選擇。

另外,上下文管理器通常和with一起使用,大大提高了程序的簡潔度,需要注意的是我們在使用with語句執行上下文操作的時候,一旦有異常拋出,異常的類型、值等拘役信息都會通過參數傳遞給__exit__()函數,我們可以自行定義相關的操作,而在對異常處理完畢以後,務必加上return True語句來保證程序的執行,否則仍然會拋出異常。

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