談一談Python的上下文管理器 – 思誠之道

發表於2016-12-28

經常在Python程式碼中看到with語句，仔細分析下，會發現這個with語句功能好強，可以自動關閉資源。這個在Python中叫上下文管理器Context Manager。那我們要怎麼用它，什麼時候用它呢。這裡我們就來聊一聊。

上下文管理器的作用

很多情況，當我們使用完一個資源後，我們需要手動的關閉掉它，比如操作檔案，建立資料庫連線等。但是，在使用資源的過程中，如果遇到異常，很可能錯誤被直接丟擲，導致來不及關閉資源。所以在大部分程式語言裡，我們使用”try-finally”語句來確保資源會關閉。比如下面的Python寫檔案程式碼：

try:
    f = open('test.txt', 'a+')
    f.write('Foo\n')
finally:
    f.close()

try:

f = open('test.txt', 'a+')

f.write('Foo\n')

finally:

f.close()

這樣做固然沒有問題，但是當”try-finally”中間的邏輯複雜，而且還帶有各種巢狀的話，程式碼就很不容易維護。Python的with語句，可以說功能同上面的”try-finally”幾乎一樣，但程式碼看上去簡潔的多，我們來實現同樣的功能：

with open('test.txt', 'a+') as f:
    f.write('Foo\n')

1 2	with open('test.txt', 'a+') as f: f.write('Foo\n')

with語句後面跟著open()方法，如果它有返回值的話，可以使用as語句將其賦值給f。在with語句塊退出時，”f.close()”方法會自動被呼叫，即使”f.write()”出現異常，也能確保close()方法被呼叫。

自定義類來使用上下文管理器

上例中”open()”方法是Python自帶的，那我們怎麼定義自己的型別來使用with語句呢。其實只要你的類定義了”__enter__()”和”__exit__()”方法，就可以使用Python的上下文管理器了。”__enter__()”方法會在with語句進入時被呼叫，其返回值會賦給as關鍵字後的變數；而”__exit__()”方法會在with語句塊退出後自動被呼叫。

我們來實現個跟上節一樣的檔案寫入功能：

class OpenFileDemo(object):
    def __init__(self, filename):
        self.filename = filename
 
    def __enter__(self):
        self.f = open(self.filename, 'a+')
        return self.f
 
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.f.close()
 
with OpenFileDemo('test.txt') as f:
    f.write('Foo\n')

class OpenFileDemo(object):

def __init__(self, filename):

self.filename = filename

def __enter__(self):

self.f = open(self.filename, 'a+')

return self.f

def __exit__(self, exc_type, exc_val, exc_tb):

self.f.close()

with OpenFileDemo('test.txt') as f:

f.write('Foo\n')

異常處理

肯定有朋友注意到上面的”__exit__()”帶了三個引數，是的，他們是用來異常處理的。大部分情況下，我們希望with語句中遇到的異常最後被丟擲，但也有時候，我們想處理這些異常。”__exit__()”方法中的三個引數exc_type, exc_val, exc_tb分別代表異常型別，異常值，和異常的Traceback。當你處理完異常後，你可以讓”__exit__()”方法返回True，此時該異常就會不會再被丟擲。比如我們將上例中的”__exit__()”方法改一下：

def __exit__(self, exc_type, exc_val, exc_tb):
        self.f.close()
        if exc_type != SyntaxError:
            return True
        return False   # Only raise exception when SyntaxError

def __exit__(self, exc_type, exc_val, exc_tb):

self.f.close()

if exc_type != SyntaxError:

return True

return False # Only raise exception when SyntaxError

現在，如果遇到SyntaxError的話，異常會被正常丟擲，而其他異常的話都會被忽略。

contextlib模組

Python中還有一個contextlib模組提供一些簡便的上下文管理器功能。

closing()方法

如果說with語句塊在退出時會自動呼叫”__exit__()”方法的話，那用了”contextlib.closing()”的with語句塊則在退出時會自動呼叫”close()”方法。看一下示例：

import contextlib
 
class Resource(object):
    def open(self):
        print 'Open Resource'
 
    def close(self):
        print 'Close Resource'
 
with contextlib.closing(Resource()) as r:
    r.open()

import contextlib

class Resource(object):

def open(self):

print 'Open Resource'

def close(self):

print 'Close Resource'

with contextlib.closing(Resource()) as r:

r.open()

程式執行後，會列印出

Open Resource
Close Resource

1 2	Open Resource Close Resource

說明Resource類建立的物件被賦給了as關鍵字後面的變數r，而with語句塊退出時，自動呼叫了”r.close()”方法。

contextmanager裝飾器

“@contextlib.contextmanager”是一個裝飾器，由它修飾的方法會有兩部分構成，中間由yield關鍵字分開。由此方法建立的上下文管理器，在程式碼塊執行前會先執行yield上面的語句；在程式碼塊執行後會再執行yield下面的語句。看個例子比較容易明白：

import contextlib
import time
 
@contextlib.contextmanager
def timeit():
    start = time.time()
    yield
    end = time.time()
    usedTime = (end - start) * 1000
    print 'Use time %d ms' % usedTime
 
with timeit():
    time.sleep(1)

import contextlib

import time

@contextlib.contextmanager

def timeit():

start = time.time()

yield

end = time.time()

usedTime = (end - start) * 1000

print 'Use time %d ms' % usedTime

with timeit():

time.sleep(1)

這個”timeit()”方法實現了一個計時器，它會計算由他生成的with語句塊執行時間。可以看出，yield上面的語句就如同之間介紹過的”__enter__()”方法，而yield下面的語句就如同”__exit__()”方法。而yield部分就是with語句塊中的程式碼。如果yield後面帶引數的話，我們就可以用as關鍵字賦值給後面的變數，比如上例：

@contextlib.contextmanager
def timeit():
    start = time.time()
    yield start
    #...
 
with timeit() as starttime:
    print starttime
    #...

@contextlib.contextmanager

def timeit():

start = time.time()

yield start

#...

with timeit() as starttime:

print starttime