Context manager nâng cao: bên trong câu lệnh with, ExitStack và các mẫu thiết kế

Bài Context Managers đã giới thiệu with, __enter__/__exit__ và @contextmanager. Bài này dành cho khi bạn cần thiết kế context manager cho thư viện của mình, hoặc gặp các tình huống khó: mở số lượng file không biết trước, khởi tạo nhiều tài nguyên mà cái sau có thể lỗi, context manager vừa dùng được như decorator, hay cần hiểu chính xác exception đi đâu.

Trong bài này, bạn sẽ học:

• Câu lệnh with thực chất được Python dịch ra thành code gì
• Ngữ nghĩa chính xác của __exit__: tham số, giá trị trả về, nuốt và thay thế exception
• @contextmanager hoạt động thế nào bên dưới và các bẫy của nó
• ExitStack: quản lý số lượng tài nguyên động, callback dọn dẹp, pop_all() để chuyển quyền sở hữu
• ContextDecorator, nullcontext, closing, chdir, suppress
• Context manager dùng một lần, dùng lại được và vào lại được (reentrant)
• Phiên bản async: AsyncExitStack, @asynccontextmanager

Câu lệnh with được dịch ra thành gì?

PEP 343 định nghĩa rằng:

with EXPR as VAR:
    BLOCK

tương đương (gần đúng) với:

manager = EXPR
enter = type(manager).__enter__       # tra trên TYPE, không phải instance
exit = type(manager).__exit__
value = enter(manager)
hit_except = False

try:
    VAR = value                        # chỉ khi có "as VAR"
    BLOCK
except BaseException as exc:
    hit_except = True
    if not exit(manager, type(exc), exc, exc.__traceback__):
        raise                          # __exit__ trả về giá trị "falsy" -> ném lại
finally:
    if not hit_except:
        exit(manager, None, None, None)

Từ đoạn code này, ta rút ra vài điều ít người để ý:

1. __enter__ và __exit__ được tra trên class, không phải trên instance. Gán obj.__exit__ = ... sẽ không có tác dụng.
2. Nếu __enter__ ném exception, __exit__ không được gọi. Tài nguyên nào đã cấp phát trong __enter__ trước khi lỗi phải tự dọn.
3. VAR nhận giá trị trả về của __enter__, không phải chính manager. Với open(), __enter__ trả về self nên hai thứ trùng nhau, nhưng với lock thì with lock as x cho x = True.
4. __exit__ được gọi với BaseException - kể cả KeyboardInterrupt, SystemExit.

Ngữ nghĩa của __exit__

def __exit__(self, exc_type, exc_value, traceback) -> bool | None:

• Không có lỗi: cả ba tham số là None.
• Có lỗi: nhận kiểu, object exception và traceback.
• Trả về truthy → exception bị nuốt, chương trình chạy tiếp sau khối with.
• Trả về falsy (kể cả None khi không có return) → exception tiếp tục lan ra ngoài.
• Nếu __exit__ tự ném exception mới, nó thay thế exception gốc (exception gốc được gắn vào __context__).

Ví dụ: một context manager chỉ nuốt những lỗi được chỉ định và ghi lại lỗi đó:

class Catch:
    def __init__(self, *exc_types):
        self.exc_types = exc_types
        self.error = None

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc, tb):
        if exc_type is not None and issubclass(exc_type, self.exc_types):
            self.error = exc
            return True            # nuốt lỗi
        return False               # lỗi khác (hoặc không lỗi) -> để nguyên

with Catch(ZeroDivisionError) as c:
    1 / 0
print("vẫn chạy tiếp, lỗi đã bắt:", repr(c.error))

with Catch(ZeroDivisionError):
    int("abc")                     # ValueError không nằm trong danh sách -> lan ra ngoài

Chú ý

Chỉ trả về True khi bạn cố ý nuốt lỗi. Một __exit__ luôn return True sẽ âm thầm nuốt cả KeyboardInterrupt - người dùng nhấn Ctrl+C mà chương trình không dừng.

Ví dụ thực tế: transaction

import sqlite3

class Transaction:
    def __init__(self, conn):
        self.conn = conn

    def __enter__(self):
        self.conn.execute("BEGIN")
        return self.conn

    def __exit__(self, exc_type, exc, tb):
        if exc_type is None:
            self.conn.execute("COMMIT")
        else:
            self.conn.execute("ROLLBACK")
        return False               # không nuốt lỗi: người gọi cần biết giao dịch thất bại

conn = sqlite3.connect(":memory:", isolation_level=None)
conn.execute("CREATE TABLE acc (name TEXT, balance INT)")
conn.execute("INSERT INTO acc VALUES ('An', 100), ('Bình', 0)")

try:
    with Transaction(conn) as c:
        c.execute("UPDATE acc SET balance = balance - 50 WHERE name = 'An'")
        raise RuntimeError("mất kết nối giữa chừng!")
        c.execute("UPDATE acc SET balance = balance + 50 WHERE name = 'Bình'")
except RuntimeError:
    pass

print(conn.execute("SELECT * FROM acc").fetchall())   # [('An', 100), ('Bình', 0)] - đã rollback

@contextmanager hoạt động thế nào?

from contextlib import contextmanager

@contextmanager
def tag(name):
    print(f"<{name}>")          # phần __enter__
    try:
        yield name              # giá trị cho "as"
    finally:
        print(f"</{name}>")     # phần __exit__

with tag("div"):
    print("nội dung")

Bên dưới, contextmanager bọc generator của bạn trong một object _GeneratorContextManager:

• __enter__ gọi next(gen) → chạy tới yield, trả về giá trị được yield.
• __exit__ khi không lỗi: gọi next(gen) lần nữa → chạy phần sau yield; generator phải kết thúc (StopIteration), nếu nó yield thêm lần nữa → RuntimeError: generator didn't stop.
• __exit__ khi có lỗi: gọi gen.throw(exc) → exception được ném vào đúng dòng yield bên trong generator.

Điều cuối cùng dẫn tới bẫy phổ biến nhất:

from contextlib import contextmanager

@contextmanager
def broken():
    print("mở tài nguyên")
    yield
    print("đóng tài nguyên")     # KHÔNG chạy nếu khối with có lỗi!

try:
    with broken():
        raise ValueError("lỗi")
except ValueError:
    pass
# In ra: "mở tài nguyên" - không bao giờ "đóng tài nguyên"

Exception được ném vào dòng yield, không có try nào bắt nên hàm dừng ngay. Luôn bọc yield trong try/finally.

Nếu muốn nuốt lỗi trong generator, bắt nó quanh yield và không ném lại:

from contextlib import contextmanager

@contextmanager
def ignore_missing():
    try:
        yield
    except FileNotFoundError:
        pass                     # nuốt lỗi -> __exit__ trả về True

with ignore_missing():
    open("không-tồn-tại.txt")
print("chạy tiếp")

Một điểm cộng ít người biết: object tạo bởi @contextmanager cũng dùng được như decorator (nó kế thừa ContextDecorator), với điều kiện mỗi lần gọi tạo generator mới:

@tag("section")
def render():
    print("thân hàm")

render()   # <section> thân hàm </section>

ExitStack: khi số lượng tài nguyên không cố định

Giả sử bạn cần mở n file mà n chỉ biết lúc chạy, và phải đảm bảo mọi file đã mở đều được đóng - kể cả khi file thứ 7 mở lỗi:

from contextlib import ExitStack

def merge_files(paths, out_path):
    with ExitStack() as stack:
        files = [stack.enter_context(open(p)) for p in paths]
        out = stack.enter_context(open(out_path, "w"))
        for f in files:
            out.write(f.read())
    # ra khỏi with: mọi file được đóng theo thứ tự NGƯỢC với lúc mở

ExitStack là một “ngăn xếp các hàm dọn dẹp”. Mỗi enter_context(cm) gọi cm.__enter__() rồi đẩy cm.__exit__ vào ngăn xếp. Khi thoát, các hàm được gọi từ trên xuống (LIFO) - và exception được truyền đúng qua từng tầng, giống như bạn viết n câu with lồng nhau.

Dọn dẹp cho những thứ không phải context manager

from contextlib import ExitStack
import shutil, tempfile

with ExitStack() as stack:
    tmp = tempfile.mkdtemp()
    stack.callback(shutil.rmtree, tmp)          # đăng ký hàm dọn dẹp bất kỳ
    stack.callback(print, "bước dọn dẹp 1")     # chạy sau cùng (LIFO)
    print("làm việc trong", tmp)

Rất hữu ích khi làm việc với API cũ chỉ có hàm open_x()/close_x().

pop_all(): khởi tạo an toàn nhiều tài nguyên trong __init__

Đây là mẫu thiết kế tinh tế nhất của ExitStack. Một class cần mở nhiều tài nguyên trong constructor. Nếu tài nguyên thứ hai lỗi, tài nguyên thứ nhất phải được đóng - nhưng nếu tất cả thành công, chúng phải sống tiếp cùng object:

from contextlib import ExitStack

class Resource:
    def __init__(self, name, fail=False):
        if fail:
            raise OSError(f"không mở được {name}")
        self.name = name
        print("mở", name)
    def __enter__(self):
        return self
    def __exit__(self, *exc):
        print("đóng", self.name)

class Service:
    def __init__(self, fail_second=False):
        with ExitStack() as stack:
            self.db = stack.enter_context(Resource("db"))
            self.cache = stack.enter_context(Resource("cache", fail=fail_second))
            # Mọi thứ thành công -> chuyển các hàm dọn dẹp sang một stack MỚI
            # stack cũ trở nên rỗng nên khi ra khỏi with, không có gì bị đóng
            self._stack = stack.pop_all()

    def close(self):
        self._stack.close()

    def __enter__(self):
        return self
    def __exit__(self, *exc):
        self.close()

with Service() as s:
    print("dùng service")
# mở db, mở cache, dùng service, đóng cache, đóng db

try:
    Service(fail_second=True)
except OSError as e:
    print("lỗi:", e)
# mở db, đóng db, lỗi: không mở được cache   <- db được dọn, không rò rỉ

Các công cụ nhỏ mà hữu ích trong contextlib

nullcontext - context manager “không làm gì”

Khi một tài nguyên là tuỳ chọn:

from contextlib import nullcontext
import threading

def process(data, lock=None):
    with lock if lock is not None else nullcontext():
        return sum(data)

print(process([1, 2, 3]))                     # không khoá
print(process([1, 2, 3], threading.Lock()))   # có khoá

Hoặc để hàm nhận cả đường dẫn lẫn file đã mở:

from contextlib import nullcontext
import io

def read_text(source):
    cm = open(source) if isinstance(source, str) else nullcontext(source)
    with cm as f:                  # nullcontext(x).__enter__() trả về x
        return f.read()

print(read_text(io.StringIO("xin chào")))   # không đóng file của người gọi

closing - cho object có .close() nhưng không phải context manager

from contextlib import closing
from urllib.request import urlopen

with closing(urlopen("https://example.com")) as page:
    ...

suppress, redirect_stdout, chdir

import os
from contextlib import suppress, redirect_stdout, chdir
import io

with suppress(FileNotFoundError):          # thay cho try/except/pass
    os.remove("file-tam.txt")

buf = io.StringIO()
with redirect_stdout(buf):                 # bắt output của code in ra màn hình
    print("bị bắt lại")
print(repr(buf.getvalue()))

with chdir("/tmp"):                        # Python 3.11+
    print(os.getcwd())                     # /tmp (hoặc /private/tmp trên macOS)

Chú ý

redirect_stdout và chdir thay đổi trạng thái toàn cục của tiến trình - không an toàn khi dùng trong code đa luồng hay async.

ContextDecorator - một class, hai cách dùng

import time
from contextlib import ContextDecorator

class timed(ContextDecorator):
    def __init__(self, label):
        self.label = label
    def __enter__(self):
        self.start = time.perf_counter()
        return self
    def __exit__(self, *exc):
        print(f"{self.label}: {time.perf_counter() - self.start:.3f}s")
        return False

with timed("khối code"):
    sum(range(1_000_000))

@timed("hàm")
def work():
    sum(range(1_000_000))

work()

Dùng một lần, dùng lại được, hay vào lại được?

Không phải context manager nào cũng dùng được nhiều lần:

Loại	Ý nghĩa	Ví dụ
Single-use	Chỉ dùng trong một câu with	file từ open(), object từ @contextmanager
Reusable	Dùng lại nhiều lần, nhưng không lồng	threading.Lock, ExitStack
Reentrant	Có thể lồng vào chính nó	threading.RLock, suppress, redirect_stdout

from contextlib import contextmanager

@contextmanager
def once():
    yield

cm = once()
with cm:
    pass
try:
    with cm:          # generator đã chạy xong, không dùng lại được
        pass
except AttributeError as e:
    print("lỗi:", e)  # '_GeneratorContextManager' object has no attribute 'args'

Thông báo lỗi khó hiểu này (CPython 3.13) chính là dấu hiệu bạn đã dùng lại một context manager single-use. Hãy gọi lại hàm (with once():) để tạo object mới mỗi lần.

Khi thiết kế context manager cho thư viện, hãy ghi rõ nó thuộc loại nào. Nếu có thể, làm cho nó reentrant hoặc báo lỗi rõ ràng khi bị dùng sai (ví dụ RuntimeError("đã được sử dụng")).

Nhiều context manager trong một câu with

Từ Python 3.10, bạn có thể dùng ngoặc để viết nhiều manager trên nhiều dòng:

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

Tương đương với hai câu with lồng nhau: nếu mở output.txt lỗi, input.txt vẫn được đóng.

Phiên bản async

Mọi thứ trên đều có phiên bản cho asyncio:

import asyncio
from contextlib import asynccontextmanager, AsyncExitStack

@asynccontextmanager
async def connection(name):
    print("kết nối", name)
    try:
        yield name
    finally:
        await asyncio.sleep(0)       # có thể await khi dọn dẹp
        print("ngắt", name)

async def main():
    async with AsyncExitStack() as stack:
        conns = [await stack.enter_async_context(connection(f"db{i}")) for i in range(3)]
        stack.push_async_callback(asyncio.sleep, 0)   # callback async
        print("dùng", conns)

asyncio.run(main())

Class tự viết thì cài đặt __aenter__ và __aexit__ (đều là async def). Xem thêm bài asyncio chuyên sâu.

Mẫu thiết kế: tạm thay đổi rồi khôi phục trạng thái

Nhiều thư viện chuẩn dùng context manager để tạm thay đổi cấu hình trong một phạm vi:

import decimal

with decimal.localcontext() as ctx:
    ctx.prec = 50                               # chỉ có hiệu lực trong khối này
    print(decimal.Decimal(1) / decimal.Decimal(7))
print(decimal.Decimal(1) / decimal.Decimal(7))  # quay về độ chính xác mặc định

Tự viết một cái tương tự cho biến môi trường:

import os
from contextlib import contextmanager

@contextmanager
def env(**overrides):
    old = {k: os.environ.get(k) for k in overrides}
    os.environ.update({k: str(v) for k, v in overrides.items()})
    try:
        yield
    finally:
        for k, v in old.items():
            if v is None:
                os.environ.pop(k, None)
            else:
                os.environ[k] = v

with env(APP_MODE="test"):
    print(os.environ["APP_MODE"])   # test
print(os.environ.get("APP_MODE"))   # None

unittest.mock.patch hoạt động đúng theo tinh thần này - vừa là context manager, vừa là decorator.

Bài tập

1. Viết context manager timeout_alarm(seconds) dùng signal.alarm (Unix) để ném TimeoutError nếu khối code chạy quá lâu, và luôn tắt alarm khi thoát.
2. Viết class TempFiles dùng ExitStack: phương thức create(suffix) tạo file tạm, mọi file bị xoá khi thoát khối with.
3. Viết @contextmanager tên atomic_write(path): ghi vào file tạm, chỉ os.replace sang path khi khối with thành công; nếu lỗi, xoá file tạm và không đụng tới file cũ.

Kết luận

• with là cú pháp gọn cho try/except/finally gọi __enter__/__exit__ tra trên class.
• __exit__ trả về truthy để nuốt lỗi - chỉ làm điều này có chủ đích.
• Trong @contextmanager, luôn bọc yield bằng try/finally; lỗi được ném vào đúng dòng yield.
• ExitStack giải quyết số lượng tài nguyên động; pop_all() giúp khởi tạo nhiều tài nguyên an toàn.
• nullcontext, closing, suppress, ContextDecorator giúp API gọn và linh hoạt.
• Ghi rõ context manager của bạn là single-use, reusable hay reentrant.

Bài tiếp theo: Descriptor - cơ chế đứng sau property, method và slots.