?凌晨三點，小陳盯著屏幕上的報錯信息，頭皮發(fā)麻。

“ResourceWarning: Unclosed file”

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

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

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

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

一、從一個沒人關的文件說起

先看一段最常見的代碼：

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

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

于是大家學會了 with 語句：

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

離開 with 代碼塊，文件自動關閉，不管里面有沒有報錯。

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

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

contextlib 就是來解決這個痛點的。

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

@contextmanager 是 contextlib 模塊里最常用、也最好用的工具。它能把一個普通的生成器函數(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)

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

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

三、一個真實的數(shù)據(jù)庫連接場景

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

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

但這樣寫有幾個隱患：如果 execute 報錯，conn.commit() 和 conn.close() 都不會執(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í)行成功還是報錯，連接都會正確關閉。成功了自動提交，失敗了自動回滾。

一個裝飾器，把資源管理的復雜度全部封裝掉了。

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

上下文管理器不只能管理“打開-關閉”類的資源，任何“進來時做一件事、出去時做另一件事”的場景都能用。

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

import time
from contextlib import contextmanager

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

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

輸出：

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

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

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

五、臨時切換目錄：用完自動恢復

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

手動寫 os.chdir 很容易忘記切回來：

import os

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

用 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ù)據(jù)')
    # 離開 with 塊自動切回原目錄

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

這個模式可以應用到很多場景：臨時修改環(huán)境變量、臨時重定向標準輸出、臨時禁用信號處理……核心思路都一樣：進去時保存狀態(tài)，出來時恢復狀態(tài)。

六、更高級的工具：ExitStack

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

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

手動處理會非常麻煩。ExitStack 就是為這種場景設計的：

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 塊時，所有文件按相反順序自動關閉

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

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

with ExitStack() as stack:
    conn = stack.callback(lambda: db.close())
    # 離開時 db.close() 會被自動調(diào)用

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

七、suppress：忽略你不關心的異常

有時候你并不想處理某個異常，只想讓它靜悄悄地過去。

比如刪除一個可能不存在的文件：

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

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

from contextlib import suppress

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

可以同時抑制多種異常：

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

八、nullcontext：需要但不需要管理的時候

寫函數(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 就是那個“什么都不做”的上下文管理器：

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 是文件對象
        # 如果 debug=False，log 是 None，但代碼仍然可以正常執(zhí)行
        if log is not None:
            log.write('處理中...')
        # 實際業(yè)務邏輯
        print("數(shù)據(jù)處理完成")

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

九、把多個上下文管理器串起來

Python 3.10 之后，contextlib 提供了一個更簡潔的寫法。以前你要嵌套多個 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())

括號把多個上下文管理器包在一起，Python 會自動按順序進入、按相反順序退出。代碼層級少了一層，看起來舒服很多。

十、實戰(zhàn)：封裝一個重試機制

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

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):
    # 這里放可能會臨時失敗的代碼
    response = requests.get('https://unstable-api.example.com/data')
    response.raise_for_status()

這個上下文管理器會自動重試 5 次，每次失敗等 2 秒。如果 5 次都失敗，拋出最后一次的異常。

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

十一、contextlib 的底層原理

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

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

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

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

寫在最后

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

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

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

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

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

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

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

最新評論