Context manager và câu lệnh with trong Python

1. Tổng quan

Mở file, kết nối database hay giữ lock đều cần một bước dọn dẹp ở cuối. Bước này phải chạy kể cả khi code ở giữa ném exception. Python giải quyết việc đó bằng context manager và câu lệnh with.

Bài này trình bày cách with hoạt động bên dưới và cách tự viết context manager bằng class hoặc bằng contextlib. Bạn cũng sẽ biết cách xử lý exception trong __exit__, cách viết test cho context manager và những lỗi hay gặp khi dùng chúng.

2. Vấn đề: dọn dẹp tài nguyên bằng try/finally

Hãy xem cách đọc file mà không dùng with:

f = open("orders.csv", encoding="utf-8")
try:
    header = f.readline()
    print(header)
finally:
    f.close()

Cách này đúng nhưng dài dòng. Nó cũng dễ quên. Nếu bạn đặt open() bên trong try, mà open() lại thất bại, thì f chưa được gán và khối finally sẽ ném NameError. Câu lệnh with gói gọn toàn bộ mẫu trên:

with open("orders.csv", encoding="utf-8") as f:
    header = f.readline()
    print(header)

print(f.closed)

order_id,customer,amount

True

Khi ra khỏi khối with, file đã được đóng. Điều này đúng cả khi thoát bình thường lẫn khi có exception.

3. Giao thức context manager

Mọi object có hai method __enter__ và __exit__ đều dùng được với with. Về mặt khái niệm, câu lệnh:

with EXPR as target:
    BODY

tương đương với đoạn code sau:

manager = EXPR
target = manager.__enter__()
try:
    BODY
except BaseException as e:
    if not manager.__exit__(type(e), e, e.__traceback__):
        raise
else:
    manager.__exit__(None, None, None)

Có ba điểm cần nhớ:

• Giá trị gán cho target là kết quả của __enter__(), không phải bản thân EXPR.
• __exit__ luôn được gọi. Khi có exception, nó nhận vào (exc_type, exc_value, traceback). Khi không có lỗi, cả ba tham số là None.
• Nếu __exit__ trả về giá trị truthy, exception bị nuốt. Ngược lại, exception được ném tiếp.

3.1. Viết context manager bằng class

Ví dụ đầu tiên là một Timer đo thời gian chạy của một khối code:

import time


class Timer:
    def __init__(self, label: str):
        self.label = label
        self.elapsed: float | None = None

    def __enter__(self):
        self._start = time.perf_counter()
        return self

    def __exit__(self, exc_type, exc, tb):
        self.elapsed = time.perf_counter() - self._start
        print(f"[{self.label}] {self.elapsed:.3f}s")
        return False


with Timer("tinh_tong_don_hang") as t:
    total = sum(range(5_000_000))

print(t.elapsed > 0)

[tinh_tong_don_hang] 0.081s
True

Giải thích từng phần:

• __enter__ ghi lại thời điểm bắt đầu và trả về self. Nhờ vậy t chính là object Timer.
• __exit__ tính thời gian, in ra màn hình rồi trả về False. Exception nếu có sẽ không bị chặn.
• Biến t vẫn dùng được sau khối with. Lý do là with không tạo scope mới.

3.2. Xử lý exception trong __exit__

Context manager phát huy tác dụng rõ nhất khi bạn cần phản ứng khác nhau tùy theo khối code thành công hay thất bại. Ví dụ dưới đây là một transaction đơn giản trên sổ cái (ledger):

class Transaction:
    def __init__(self, ledger: list):
        self.ledger = ledger
        self.pending: list = []

    def __enter__(self):
        return self.pending

    def __exit__(self, exc_type, exc, tb):
        if exc_type is None:
            self.ledger.extend(self.pending)
            print(f"commit {len(self.pending)} bản ghi")
        else:
            print(f"rollback do {exc_type.__name__}: {exc}")
        return False


ledger = []

with Transaction(ledger) as tx:
    tx.append({"order_id": 1, "amount": 250_000})
    tx.append({"order_id": 2, "amount": 120_000})

try:
    with Transaction(ledger) as tx:
        tx.append({"order_id": 3, "amount": 90_000})
        raise ValueError("số dư không đủ")
except ValueError:
    pass

print(len(ledger))

commit 2 bản ghi
rollback do ValueError: số dư không đủ
2

Ở khối with thứ hai, đơn hàng số 3 không được ghi vào ledger. Vì __exit__ trả về False, ValueError vẫn được ném ra ngoài và người gọi tự quyết định cách xử lý.

Lưu ý: Chỉ trả về True từ __exit__ khi bạn cố ý muốn nuốt exception. Một method __exit__ vô tình return self hoặc trả về một object truthy khác sẽ làm lỗi biến mất mà không có dấu vết nào.

4. Viết gọn hơn với @contextmanager

Viết cả một class chỉ cho một bước setup và teardown thì khá dài. Module contextlib cung cấp decorator @contextmanager để biến một generator thành context manager:

from contextlib import contextmanager


class PaymentService:
    def connect(self):
        print("mở kết nối cổng thanh toán")

    def charge(self, customer: str, amount: int):
        print(f"trừ {amount} của {customer}")

    def disconnect(self):
        print("đóng kết nối")


@contextmanager
def payment_session(service: PaymentService):
    service.connect()
    try:
        yield service
    finally:
        service.disconnect()


with payment_session(PaymentService()) as payment:
    payment.charge("Lan", 300_000)

mở kết nối cổng thanh toán
trừ 300000 của Lan
đóng kết nối

Cách decorator ánh xạ generator sang giao thức context manager:

• Code trước yield đóng vai trò __enter__.
• Giá trị được yield là giá trị gán cho biến sau as.
• Code sau yield đóng vai trò __exit__. Nếu khối with ném exception, exception đó được ném lại tại vị trí yield.

Vì exception xuất hiện ngay tại yield, khối try/finally quanh yield là bắt buộc. Thiếu nó thì disconnect() sẽ không chạy khi có lỗi.

Muốn xử lý riêng từng loại lỗi, bạn dùng thêm except:

@contextmanager
def safe_charge(order_id: int):
    try:
        yield
    except ConnectionError as e:
        print(f"đơn {order_id}: lỗi mạng, sẽ thử lại sau ({e})")


with safe_charge(42):
    raise ConnectionError("timeout")

print("tiếp tục xử lý")

đơn 42: lỗi mạng, sẽ thử lại sau (timeout)
tiếp tục xử lý

Generator bắt lỗi nhưng không ném lại, nên exception bị nuốt. Kết quả này tương đương với việc __exit__ trả về True.

5. Các tiện ích có sẵn trong contextlib

5.1. suppress: bỏ qua một loại exception cụ thể

import os
from contextlib import suppress

with suppress(FileNotFoundError):
    os.remove("cache_orders.tmp")

print("xong")

xong

Đoạn trên ngắn hơn và nói rõ ý đồ hơn so với try: ... except FileNotFoundError: pass.

5.2. ExitStack: số lượng tài nguyên không biết trước

Giả sử bạn cần gộp nhiều file hóa đơn, nhưng số file chỉ biết lúc runtime:

from contextlib import ExitStack
from pathlib import Path


def merge_invoices(paths: list[Path], output: Path) -> None:
    with ExitStack() as stack:
        files = [stack.enter_context(open(p, encoding="utf-8")) for p in paths]
        with open(output, "w", encoding="utf-8") as out:
            for f in files:
                out.write(f.read())


for i in range(3):
    Path(f"invoice_{i}.txt").write_text(f"hoa don {i}\n", encoding="utf-8")

merge_invoices([Path(f"invoice_{i}.txt") for i in range(3)], Path("all.txt"))
print(Path("all.txt").read_text(encoding="utf-8"))

hoa don 0
hoa don 1
hoa don 2

enter_context gọi __enter__ và ghi nhớ __exit__ tương ứng. Khi ExitStack kết thúc, các file được đóng theo thứ tự ngược lại lúc mở. Nếu file thứ ba mở thất bại, hai file đầu vẫn được đóng đúng cách.

5.3. closing và nullcontext

• closing(obj) biến một object chỉ có method close() thành context manager.
• nullcontext(value) là context manager không làm gì cả. Nó hữu ích khi một tham số có thể là context manager hoặc không:

from contextlib import nullcontext
import threading

def update_stock(product: dict, qty: int, lock: threading.Lock | None = None):
    with lock if lock is not None else nullcontext():
        product["stock"] -= qty

product = {"sku": "AO-01", "stock": 10}
update_stock(product, 3)
update_stock(product, 2, threading.Lock())
print(product["stock"])

5

6. Nhiều context manager trong một câu lệnh with

Bạn có thể khai báo nhiều context manager trong cùng một câu lệnh with. Từ Python 3.10, cú pháp dùng ngoặc đơn cho phép xuống dòng gọn gàng:

with (
    open("orders.csv", encoding="utf-8") as src,
    open("orders_backup.csv", "w", encoding="utf-8") as dst,
):
    dst.write(src.read())

Cách viết này tương đương với hai khối with lồng nhau. src được mở trước và đóng sau cùng.

7. Async context manager

Với code asyncio, giao thức tương ứng là __aenter__ và __aexit__, dùng cùng async with. Phiên bản generator của nó là @asynccontextmanager:

import asyncio
from contextlib import asynccontextmanager


@asynccontextmanager
async def order_queue(name: str):
    print(f"kết nối queue {name}")
    await asyncio.sleep(0.01)
    try:
        yield [f"order-{i}" for i in range(3)]
    finally:
        await asyncio.sleep(0.01)
        print(f"đóng queue {name}")


async def main():
    async with order_queue("payments") as messages:
        for m in messages:
            print("xử lý", m)

asyncio.run(main())

kết nối queue payments
xử lý order-0
xử lý order-1
xử lý order-2
đóng queue payments

Lưu ý: Dùng with thường với một async context manager sẽ gây lỗi TypeError (hoặc AttributeError ở các bản Python cũ). Ngược lại cũng vậy. Hai giao thức này tách biệt nhau.

8. Kiểm thử context manager

Context manager nên được test cả nhánh thành công lẫn nhánh lỗi. Bạn cài pytest trước:

pip install pytest

Đặt class Transaction ở mục 3.2 vào file transaction.py, rồi viết test:

# test_transaction.py
import pytest
from transaction import Transaction


def test_commit_when_no_error():
    ledger = []
    with Transaction(ledger) as tx:
        tx.append({"order_id": 1, "amount": 100})

    assert ledger == [{"order_id": 1, "amount": 100}]


def test_rollback_and_propagate_error():
    ledger = []
    with pytest.raises(ValueError, match="số dư"):
        with Transaction(ledger) as tx:
            tx.append({"order_id": 2, "amount": 999})
            raise ValueError("số dư không đủ")

    assert ledger == []

pytest -q

..                                                       [100%]
2 passed in 0.01s

Test thứ hai kiểm tra hai điều cùng lúc: dữ liệu bị rollback, và exception không bị nuốt. Bản thân pytest.raises cũng là một context manager, một ví dụ quen thuộc của chính chủ đề bài này.

9. Lỗi hay gặp và khi nào không nên dùng

Quên try/finally quanh yield. Đây là lỗi phổ biến nhất với @contextmanager:

@contextmanager
def broken_session(service):
    service.connect()
    yield service
    service.disconnect()  # không chạy nếu khối with ném exception

Dùng tài nguyên sau khi with kết thúc. Biến vẫn còn tồn tại, nhưng tài nguyên đã được giải phóng:

with open("orders.csv", encoding="utf-8") as f:
    pass
f.read()

ValueError: I/O operation on closed file.

yield nhiều lần trong generator. Generator dùng với @contextmanager phải yield đúng một lần. Nếu yield nhiều hơn, Python ném RuntimeError: generator didn't stop.

Tái sử dụng context manager tạo từ @contextmanager. Object này chỉ dùng được một lần. Gọi with lần hai trên cùng object sẽ lỗi. Mỗi lần dùng, bạn hãy gọi lại hàm để tạo object mới.

Lạm dụng. Không nên dùng context manager trong các trường hợp sau:

• Tài nguyên cần sống lâu hơn một khối code, ví dụ connection pool tồn tại suốt vòng đời ứng dụng. Khi đó hãy quản lý nó ở tầng khởi động và tắt ứng dụng.
• Logic không có cặp “mở/đóng” rõ ràng. Một hàm thông thường sẽ dễ đọc hơn.
• Muốn bỏ qua mọi exception bằng suppress(Exception). Cách này che giấu bug giống hệt except: pass.

Lưu ý: Đừng dựa vào garbage collector hay __del__ để đóng tài nguyên. Thời điểm chúng chạy không được đảm bảo, đặc biệt trên các implementation như PyPy. Hãy dùng with.

10. Kết luận

Context manager gói phần setup và teardown vào một chỗ, và đảm bảo teardown luôn chạy kể cả khi có exception. Bạn đã thấy:

• Giao thức __enter__ và __exit__, cùng ý nghĩa của giá trị trả về từ __exit__.
• Cách viết gọn bằng @contextmanager, với try/finally quanh yield là bắt buộc.
• Các tiện ích suppress, ExitStack, closing, nullcontext, cùng phiên bản async.
• Cách test cả nhánh commit lẫn nhánh rollback bằng pytest.

Khi nào code của bạn có cặp “lấy rồi trả” (file, lock, transaction, kết nối), hãy cân nhắc bọc nó trong một context manager.