?凌晨三點(diǎn)，小陳盯著屏幕上的報(bào)錯(cuò)信息，頭皮發(fā)麻。

“ResourceWarning: Unclosed file”

就這一行警告，讓他在一堆歷史代碼里翻了兩個(gè)小時(shí)。打開的文件忘記關(guān)了，數(shù)據(jù)庫連接沒釋放，臨時(shí)修改的目錄路徑也沒改回來。代碼跑起來沒問題，但跑久了服務(wù)器就開始報(bào)“too many open files”。

同事老張路過，瞥了一眼屏幕：“你還在手動(dòng)寫 try...finally 呢？用 contextlib 啊，幾行裝飾器的事。”

小陳一臉懵：“那是什么？”

這就是很多 Python 開發(fā)者都會(huì)經(jīng)歷的階段——被資源管理問題折磨過后，才發(fā)現(xiàn)標(biāo)準(zhǔn)庫里藏著一個(gè)寶藏模塊。

一、從一個(gè)沒人關(guān)的文件說起

先看一段最常見的代碼：

f = open('data.txt', 'r')
data = f.read()
print(data)
f.close()

寫過 Python 的人都知道，這樣寫不夠安全。如果 f.read() 中間拋異常，f.close() 永遠(yuǎn)不會(huì)執(zhí)行，文件句柄就泄露了。

于是大家學(xué)會(huì)了 with 語句：

with open('data.txt', 'r') as f:
    data = f.read()
    print(data)

離開 with 代碼塊，文件自動(dòng)關(guān)閉，不管里面有沒有報(bào)錯(cuò)。

這個(gè) with 語句背后的原理，就是上下文管理器。一個(gè)類只要實(shí)現(xiàn)了 __enter__ 和 __exit__ 兩個(gè)魔法方法，就能放進(jìn) with 里用。

但問題來了：每次都要寫一個(gè)完整的類，就為了管理一個(gè)資源？太啰嗦了。

contextlib 就是來解決這個(gè)痛點(diǎn)的。

二、contextlib 的核心武器：@contextmanager 裝飾器

@contextmanager 是 contextlib 模塊里最常用、也最好用的工具。它能把一個(gè)普通的生成器函數(shù)直接變成上下文管理器。

看看怎么用：

from contextlib import contextmanager

@contextmanager
def managed_file(filename):
    f = open(filename, 'r')
    try:
        yield f
    finally:
        f.close()

# 使用方式
with managed_file('data.txt') as f:
    data = f.read()
    print(data)

關(guān)鍵點(diǎn)在這里： yield 前面的代碼相當(dāng)于 __enter__，yield 后面的代碼（放在 finally 里）相當(dāng)于 __exit__。不管 with 代碼塊里發(fā)生什么，finally 都會(huì)執(zhí)行，文件一定能關(guān)掉。

這個(gè)寫法比寫一個(gè)完整的類清爽太多了。幾行代碼搞定，邏輯一目了然。

三、一個(gè)真實(shí)的數(shù)據(jù)庫連接場(chǎng)景

假設(shè)你在寫一個(gè) Web 爬蟲，需要把數(shù)據(jù)存到 SQLite 數(shù)據(jù)庫。每次操作都要打開連接、獲取游標(biāo)、提交事務(wù)、關(guān)閉連接，寫起來非常繁瑣：

def save_data(data):
    conn = sqlite3.connect('app.db')
    cursor = conn.cursor()
    cursor.execute('INSERT INTO items VALUES (?)', (data,))
    conn.commit()
    conn.close()

但這樣寫有幾個(gè)隱患：如果 execute 報(bào)錯(cuò)，conn.commit() 和 conn.close() 都不會(huì)執(zhí)行，數(shù)據(jù)庫連接就懸在那里了。

用 @contextmanager 包裝一下：

@contextmanager
def get_db():
    conn = sqlite3.connect('app.db')
    try:
        yield conn
        conn.commit()
    except Exception:
        conn.rollback()
        raise
    finally:
        conn.close()

# 使用
with get_db() as conn:
    cursor = conn.cursor()
    cursor.execute('INSERT INTO items VALUES (?)', ('test',))

現(xiàn)在不管代碼執(zhí)行成功還是報(bào)錯(cuò)，連接都會(huì)正確關(guān)閉。成功了自動(dòng)提交，失敗了自動(dòng)回滾。

一個(gè)裝飾器，把資源管理的復(fù)雜度全部封裝掉了。

四、計(jì)時(shí)器：用上下文做性能監(jiān)控

上下文管理器不只能管理“打開-關(guān)閉”類的資源，任何“進(jìn)來時(shí)做一件事、出去時(shí)做另一件事”的場(chǎng)景都能用。

比如你想測(cè)一段代碼的執(zhí)行時(shí)間：

import time
from contextlib import contextmanager

@contextmanager
def timer(name):
    start = time.time()
    print(f"{name} 開始...")
    yield
    elapsed = time.time() - start
    print(f"{name} 完成，耗時(shí) {elapsed:.2f} 秒")

# 使用
with timer("數(shù)據(jù)清洗"):
    # 這里放你要測(cè)的代碼
    data = [i**2 for i in range(1000000)]
    print(f"生成了 {len(data)} 條數(shù)據(jù)")

輸出：

數(shù)據(jù)清洗 開始...
生成了 1000000 條數(shù)據(jù)
數(shù)據(jù)清洗 完成，耗時(shí) 0.18 秒

不需要寫一堆 start = time.time() 和 print(...) 的重復(fù)代碼。把計(jì)時(shí)邏輯包進(jìn)上下文管理器里，用的時(shí)候一行 with timer(...) 就搞定了。

這對(duì)性能調(diào)優(yōu)特別有用。 你可以快速給多個(gè)代碼塊加上計(jì)時(shí)，找出瓶頸在哪里。

五、臨時(shí)切換目錄：用完自動(dòng)恢復(fù)

寫腳本的時(shí)候，經(jīng)常需要臨時(shí)切換工作目錄去處理文件。處理完得切回來，不然會(huì)影響后面的代碼。

手動(dòng)寫 os.chdir 很容易忘記切回來：

import os

os.chdir('/tmp')
# 處理臨時(shí)文件...
# 糟糕，忘記切回原來的目錄了
os.remove('important.txt')  # 刪錯(cuò)了！

用 contextlib 包裝一下：

@contextmanager
def cd(path):
    old_dir = os.getcwd()
    os.chdir(path)
    try:
        yield
    finally:
        os.chdir(old_dir)

# 使用
with cd('/tmp'):
    # 在 /tmp 目錄下操作
    with open('temp.txt', 'w') as f:
        f.write('臨時(shí)數(shù)據(jù)')
    # 離開 with 塊自動(dòng)切回原目錄

# 這里已經(jīng)回到原來的目錄了

這個(gè)模式可以應(yīng)用到很多場(chǎng)景：臨時(shí)修改環(huán)境變量、臨時(shí)重定向標(biāo)準(zhǔn)輸出、臨時(shí)禁用信號(hào)處理……核心思路都一樣：進(jìn)去時(shí)保存狀態(tài)，出來時(shí)恢復(fù)狀態(tài)。

六、更高級(jí)的工具：ExitStack

有時(shí)候你需要同時(shí)管理多個(gè)資源，而且這些資源的數(shù)量在運(yùn)行時(shí)才能確定。比如打開一批文件：

files = []
for filename in file_list:
    files.append(open(filename, 'r'))
# 萬一中間某個(gè)文件打開失敗，前面已經(jīng)打開的文件怎么關(guān)？

手動(dòng)處理會(huì)非常麻煩。ExitStack 就是為這種場(chǎng)景設(shè)計(jì)的：

from contextlib import ExitStack

with ExitStack() as stack:
    files = []
    for filename in file_list:
        f = stack.enter_context(open(filename, 'r'))
        files.append(f)
    # 所有文件都成功打開，繼續(xù)處理
    for f in files:
        print(f.read())
# 離開 with 塊時(shí)，所有文件按相反順序自動(dòng)關(guān)閉

ExitStack 內(nèi)部維護(hù)了一個(gè)棧。每次調(diào)用 enter_context，它就把這個(gè)資源記下來。離開 with 塊時(shí)，按照后進(jìn)先出的順序自動(dòng)清理所有資源。不管中間哪個(gè)步驟出問題，已經(jīng)成功打開的資源都會(huì)被正確關(guān)閉。

還有一個(gè)常用場(chǎng)景：在舊代碼里，有些資源不是上下文管理器，只有 close 方法。ExitStack 也能處理：

with ExitStack() as stack:
    conn = stack.callback(lambda: db.close())
    # 離開時(shí) db.close() 會(huì)被自動(dòng)調(diào)用

callback 方法讓你可以注冊(cè)任意的清理函數(shù)，非常靈活。

七、suppress：忽略你不關(guān)心的異常

有時(shí)候你并不想處理某個(gè)異常，只想讓它靜悄悄地過去。

比如刪除一個(gè)可能不存在的文件：

try:
    os.remove('temp.txt')
except FileNotFoundError:
    pass

寫 try-except-pass 太啰嗦了。contextlib.suppress 專門解決這個(gè)問題：

from contextlib import suppress

with suppress(FileNotFoundError):
    os.remove('temp.txt')

可以同時(shí)抑制多種異常：

代碼干凈了很多，意圖也很明確：“這個(gè)異常出現(xiàn)了也沒關(guān)系，忽略它。”

八、nullcontext：需要但不需要管理的時(shí)候

寫函數(shù)的時(shí)候，有時(shí)候需要根據(jù)參數(shù)決定是否使用上下文管理器。比如調(diào)試模式下打開日志文件，生產(chǎn)模式下什么都不做：

def process_data(debug=False):
    if debug:
        manager = open('debug.log', 'w')
    else:
        manager = ???  # 這里放什么？

    with manager as log:
        log.write('處理中...')

nullcontext 就是那個(gè)“什么都不做”的上下文管理器：

from contextlib import nullcontext

def process_data(debug=False):
    if debug:
        manager = open('debug.log', 'w')
    else:
        manager = nullcontext()

    with manager as log:
        # 如果 debug=True，log 是文件對(duì)象
        # 如果 debug=False，log 是 None，但代碼仍然可以正常執(zhí)行
        if log is not None:
            log.write('處理中...')
        # 實(shí)際業(yè)務(wù)邏輯
        print("數(shù)據(jù)處理完成")

這樣你就不需要寫兩套邏輯，一套帶 with 一套不帶。統(tǒng)一用 with 結(jié)構(gòu)，nullcontext 會(huì)乖乖地什么也不做。

九、把多個(gè)上下文管理器串起來

Python 3.10 之后，contextlib 提供了一個(gè)更簡(jiǎn)潔的寫法。以前你要嵌套多個(gè) with：

with open('input.txt') as infile:
    with open('output.txt', 'w') as outfile:
        outfile.write(infile.read())

現(xiàn)在可以寫成一行：

with (
    open('input.txt') as infile,
    open('output.txt', 'w') as outfile
):
    outfile.write(infile.read())

括號(hào)把多個(gè)上下文管理器包在一起，Python 會(huì)自動(dòng)按順序進(jìn)入、按相反順序退出。代碼層級(jí)少了一層，看起來舒服很多。

十、實(shí)戰(zhàn)：封裝一個(gè)重試機(jī)制

把這些技巧組合起來，能做出很實(shí)用的工具。比如一個(gè)帶重試功能的上下文管理器：

import time
from contextlib import contextmanager

@contextmanager
def retry(max_attempts=3, delay=1):
    last_exception = None
    for attempt in range(max_attempts):
        try:
            yield
            return  # 成功就退出
        except Exception as e:
            last_exception = e
            print(f"第 {attempt + 1} 次嘗試失敗: {e}")
            if attempt < max_attempts - 1:
                time.sleep(delay)
    raise last_exception

# 使用
with retry(max_attempts=5, delay=2):
    # 這里放可能會(huì)臨時(shí)失敗的代碼
    response = requests.get('https://unstable-api.example.com/data')
    response.raise_for_status()

這個(gè)上下文管理器會(huì)自動(dòng)重試 5 次，每次失敗等 2 秒。如果 5 次都失敗，拋出最后一次的異常。

調(diào)用方的代碼非常干凈，不需要寫任何重試邏輯。這就是 contextlib 的魅力——把橫切關(guān)注點(diǎn)（cross-cutting concerns）封裝起來，讓業(yè)務(wù)代碼保持簡(jiǎn)潔。

十一、contextlib 的底層原理

用 @contextmanager 裝飾一個(gè)生成器函數(shù)，Python 背后做了這些事：

• 調(diào)用函數(shù)，得到一個(gè)生成器對(duì)象
• 調(diào)用生成器的 __next__()，執(zhí)行到 yield 暫停，yield 的值作為 __enter__ 的返回值
• 執(zhí)行 with 代碼塊
• 代碼塊正常結(jié)束或拋出異常時(shí)，調(diào)用生成器的 throw() 或 __next__()，執(zhí)行 yield 后面的代碼
• 生成器執(zhí)行結(jié)束

所以 yield 后面的代碼一定要放在 try-finally 里，確保不管 with 塊里發(fā)生了什么，清理邏輯都能執(zhí)行。

這個(gè)設(shè)計(jì)非常巧妙。 生成器本來是用來產(chǎn)生序列的，Python 把它復(fù)用到上下文管理器的場(chǎng)景，用 yield 切開了“進(jìn)入”和“退出”兩個(gè)階段。

寫在最后

回頭再看小陳的故事。如果當(dāng)時(shí)他知道 @contextmanager，打開數(shù)據(jù)庫連接的那段代碼會(huì)寫成這樣：

@contextmanager
def get_conn():
    conn = create_conn()
    try:
        yield conn
    finally:
        conn.close()

三個(gè)函數(shù)調(diào)用，一個(gè) yield，一個(gè) finally，搞定。不用寫類，不用記 __enter__ 和 __exit__，代碼意圖清晰得不能再清晰。

contextlib 這個(gè)模塊不大，但每一行代碼都經(jīng)過精心設(shè)計(jì)。 它解決的是一個(gè)很具體的問題——讓 with 語句的編寫變得簡(jiǎn)單。但因?yàn)檫@個(gè)“具體問題”在編程里幾乎天天遇到，它的價(jià)值就被無限放大了。

下次你再遇到需要“進(jìn)去時(shí)做點(diǎn)事、出來時(shí)做點(diǎn)事”的場(chǎng)景，先想想能不能用 @contextmanager 包裝一下。寫出一個(gè)漂亮的上下文管理器，那種“用起來真舒服”的感覺，比寫一百行注釋都來得實(shí)在。

彩蛋： contextlib 還有 AbstractContextManager、AsyncContextDecorator 等高級(jí)工具，適合在寫框架或庫的時(shí)候使用。不過對(duì)于 90% 的日常開發(fā)，@contextmanager 加上 ExitStack、suppress、nullcontext 這四個(gè)工具，已經(jīng)足夠應(yīng)對(duì)絕大多數(shù)資源管理問題了。

以上就是深入淺出Python contextlib如何優(yōu)雅管理上下文資源的詳細(xì)內(nèi)容，更多關(guān)于Python contextlib管理上下文的資料請(qǐng)關(guān)注腳本之家其它相關(guān)文章！

最新評(píng)論