Context manager trong Python: giải thích câu lệnh with

Câu lệnh tự dọn dẹp sau khi dùng xong

Mọi tài nguyên bạn mở trong chương trình (một tệp, một kết nối mạng, một handle cơ sở dữ liệu, một khóa) đều cần được đóng khi dùng xong. Nếu quên, bạn sẽ rò rỉ bộ nhớ, giữ những khóa chặn các tiến trình khác, hoặc làm hỏng tệp khi chương trình bị sập. Câu lệnh with của Python lo việc này giúp bạn.

Mẫu mà ai cũng bắt đầu là đọc một tệp:

with open("notes.txt") as f:
    contents = f.read()
    print(contents)

Có hai việc diễn ra tự động. Khi vào khối, open() trả về một đối tượng tệp gán cho f. Khi ra khỏi khối, dù khối chạy xong bình thường, trả về sớm hay ném ra ngoại lệ, Python gọi f.close() giúp bạn.

Vậy thôi. Đó là toàn bộ ý nghĩa của with.

with thay thế cho điều gì

Trước khi có context manager, đoạn code an toàn tương đương là một khối try/finally:

f = open("notes.txt")
try:
    contents = f.read()
    print(contents)
finally:
    f.close()

Năm dòng nghi thức chỉ để "đọc một tệp rồi đóng nó khi xong". Nhân con số đó lên với mỗi lần gọi open trong một chương trình lớn, bạn sẽ thấy ngay sức hấp dẫn của with. with ngắn hơn, khó viết sai hơn, và không thể quên bước dọn dẹp.

Mở nhiều tài nguyên

Bạn có thể khai báo nhiều context manager trong một câu lệnh with:

with open("input.txt") as src, open("output.txt", "w") as dst:
    dst.write(src.read().upper())

Cả hai tệp được mở khi vào khối và được đóng khi ra khỏi khối. Nếu lệnh open thứ nhất thành công còn lệnh thứ hai ném ra ngoại lệ, Python vẫn đóng tệp thứ nhất: cơ chế này xử lý đúng cả trường hợp mới thiết lập được một phần.

Với danh sách tài nguyên dài hơn, cách viết trong ngoặc đơn (Python 3.10 trở lên) rõ ràng hơn:

with (
    open("a.txt") as a,
    open("b.txt") as b,
    open("c.txt") as c,
):
    ...

Context manager thật ra là gì

Bất kỳ đối tượng nào định nghĩa __enter__ và __exit__ đều là một context manager. Giao thức này cực kỳ đơn giản:

• __enter__(self) chạy khi khối with bắt đầu. Giá trị nó trả về chính là thứ được gán cho as name.
• __exit__(self, exc_type, exc_value, traceback) chạy khi khối kết thúc, bất kể kết thúc thế nào. Nếu một ngoại lệ gây ra việc thoát, thông tin ngoại lệ được truyền vào để context manager có thể xem xét hoặc chặn nó lại.

Đây là một context manager tối giản, đo thời gian của khối lệnh mà nó bao quanh:

import time

class Timer:
    def __enter__(self):
        self.start = time.perf_counter()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        elapsed = time.perf_counter() - self.start
        print(f"Elapsed: {elapsed*1000:.2f} ms")

with Timer():
    total = sum(i * i for i in range(1_000_000))
    print(total)

with Timer(): tạo đối tượng, gọi __enter__ của nó, chạy phần thân, rồi gọi __exit__. Không tệp, không khóa: chỉ là một lớp bọc nhỏ quanh việc "làm gì đó, đo xem mất bao lâu".

Lối tắt contextlib.contextmanager

Định nghĩa một lớp cho mỗi context manager thì nặng nề hơn mức cần thiết. contextlib.contextmanager biến một hàm generator thành context manager: một lệnh yield chia phần "trước" và phần "sau":

from contextlib import contextmanager
import time

@contextmanager
def timed(label):
    start = time.perf_counter()
    try:
        yield
    finally:
        elapsed = time.perf_counter() - start
        print(f"{label}: {elapsed*1000:.2f} ms")

with timed("big sum"):
    total = sum(i * i for i in range(1_000_000))
    print(total)

Mọi thứ trước yield là hành vi của __enter__. Mọi thứ sau đó là __exit__. Khối try/finally đảm bảo phần dọn dẹp vẫn chạy kể cả khi phần thân ném ra ngoại lệ.

Phần lớn context manager bạn tự viết đều có dạng này. Hãy dùng dạng decorator trước; chỉ chuyển sang lớp khi bạn cần điều gì đó mà dạng generator không diễn đạt được.

Tạm thời thay đổi một thứ

Một dạng rất hay gặp: đặt một giá trị, dùng nó, rồi khôi phục lại. Context manager diễn đạt điều này rất gọn:

from contextlib import contextmanager

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

# Run some code with a different working directory,
# then snap back automatically when the block ends.
with cd("/tmp"):
    import os
    print("cwd:", os.getcwd())

import os
print("back to:", os.getcwd())

Mọi mẫu "đặt rồi khôi phục" (biến môi trường, mức chi tiết của log, cờ tính năng, fixture khi kiểm thử) đều tự nhiên khớp với context manager. Người gọi không phải nhớ khôi phục gì cả.

Chặn ngoại lệ

Phương thức __exit__ có thể trả về True để báo với Python "tôi đã xử lý ngoại lệ, hãy bỏ qua nó". Việc này hiếm gặp và thường là dấu hiệu code có vấn đề, nhưng đó chính là cách contextlib.suppress hoạt động:

from contextlib import suppress

with suppress(FileNotFoundError):
    with open("maybe.txt") as f:
        print(f.read())

print("carried on")

suppress(FileNotFoundError) biến FileNotFoundError thành một thao tác không làm gì. Hãy dùng nó cho những việc thật sự không bắt buộc, kiểu "thử cái này, không chạy được cũng không sao". Đừng dùng nó để làm im những ngoại lệ mà bạn chưa nghĩ kỹ.

Những context manager khác bạn sẽ gặp

Một khi để ý, bạn sẽ thấy context manager xuất hiện khắp thư viện chuẩn:

import threading
from pathlib import Path

# Locks - guarantee release even if the critical section raises.
lock = threading.Lock()
with lock:
    ...

# tempfile - delete the temp file when you're done.
from tempfile import TemporaryDirectory
with TemporaryDirectory() as tmp:
    path = Path(tmp) / "scratch.txt"
    path.write_text("hello")

# Database connections - close the connection (or end the transaction).
import sqlite3
with sqlite3.connect(":memory:") as conn:
    conn.execute("CREATE TABLE t (x INTEGER)")

Các thư viện bên thứ ba cũng theo đúng quy ước này. Khi bạn thấy with something as x:, gần như lúc nào nó cũng có nghĩa là "dùng x trong suốt khối này, rồi dọn dẹp sau đó".

Khi nào không nên dùng with

• Khi thật ra không có bước thiết lập và dọn dẹp. Bọc một đoạn code bất kỳ trong context manager mà không có lý do chỉ thêm rối.
• Khi bạn cần tài nguyên xuyên suốt nhiều khối lệnh không liên quan. Giữ một with mở suốt vòng đời của một script dài có thể che mất phạm vi dọn dẹp thực sự. Hãy cân nhắc một lớp sở hữu tài nguyên đó.
• Khi decorator hợp hơn. Một số mẫu lặp lại (thử lại, ghi log, đo thời gian) đọc tự nhiên hơn dưới dạng @decorator đặt trên hàm thay vì with ...: bên trong hàm. Hãy chọn cách đọc dễ hơn tại chỗ gọi.

Phần lớn thời gian, with là lựa chọn đúng. Những ngoại lệ hiếm hoi rất dễ nhận ra khi bạn để ý.

Tiếp theo: làm việc với tệp thật

Giờ bạn đã biết cơ chế đằng sau with open(...) as f:, cũng là ngữ cảnh mà bạn sẽ dùng nó trong chín mươi phần trăm trường hợp. Chương tiếp theo sẽ dùng nó để đọc, ghi và duyệt các tệp trên ổ đĩa.