Làm chủ contextlib: Bí kíp quản lý tài nguyên trong Python không lo rò rỉ bộ nhớ

Tại sao bạn nên quan tâm đến contextlib?

Nếu bạn từng viết Python, chắc hẳn bạn đã quen với with open('data.txt') as f:. Đây chính là Context Manager. Nó giúp bạn mở file và tự động đóng lại mà không cần bận tâm, giúp tránh rò rỉ bộ nhớ hay treo kết nối.

Hồi mới làm automation, mình hay viết code kiểu mở file rồi định bụng cuối hàm mới close(). Kết quả là một lần chạy script đẩy log liên tục, mình quên đóng socket trong vòng lặp. Chỉ sau 4 tiếng, server báo lỗi “Too many open files” và sập toàn bộ service. Bài học xương máu đó khiến mình nhận ra: Quản lý tài nguyên thủ công là tự sát.

Thông thường, tạo Context Manager đòi hỏi bạn viết một class với __enter__ và __exit__ khá rườm rà. contextlib ra đời để dẹp bỏ sự cồng kềnh đó. Nó cung cấp các công cụ để bạn tạo trình quản lý ngữ cảnh chỉ với vài dòng code. Code của bạn sẽ vừa sạch, vừa đậm chất “Pythonic”.

Mình thường dùng contextlib để xử lý các task như tạm thời đổi thư mục làm việc hoặc tắt log khi chạy test. Nó cũng đảm bảo các kết nối database luôn được trả về pool, ngay cả khi script gặp sự cố giữa chừng.

Tạo Context Manager siêu tốc với @contextmanager

Tin vui là contextlib nằm sẵn trong thư viện chuẩn của Python. Bạn chỉ cần import là dùng ngay, không cần pip install rắc rối. Công cụ đáng giá nhất ở đây là decorator @contextmanager.

Thay vì tạo class, bạn chỉ cần biến một hàm generator thành Context Manager. Hãy xem cách mình tạo một công cụ đo thời gian thực thi cực kỳ tiện lợi dưới đây:

import time
from contextlib import contextmanager

@contextmanager
def execution_timer(label):
    start = time.perf_counter()
    try:
        # Code trong khối 'with' sẽ chạy tại đây
        yield
    finally:
        end = time.perf_counter()
        print(f"[{label}] Hoàn thành sau: {end - start:.4f} giây")

# Sử dụng thực tế
with execution_timer("Tính toán dữ liệu 10 triệu dòng"):
    time.sleep(1.2)  # Giả lập tác vụ nặng
    print("Đang xử lý...")

Trong ví dụ này, phần trước yield đóng vai trò chuẩn bị (setup), còn phần trong finally là dọn dẹp (cleanup). Việc đặt yield trong khối try...finally là bắt buộc. Nó đảm bảo đoạn code cleanup luôn chạy, dù code bên trong with có bị crash hay không.

3 “Vũ khí” bí mật giúp code gọn gàng hơn

Ít người biết rằng contextlib còn sở hữu những hàm tiện ích cực hay, giúp loại bỏ các khối code lặp đi lặp lại.

1. contextlib.suppress: Lờ đi các lỗi không quan trọng

Đôi khi bạn muốn xóa một file tạm nhưng không muốn script dừng lại nếu file đó không tồn tại. Thay vì dùng try...except FileNotFoundError: pass dài dòng, hãy dùng suppress.

import os
from contextlib import suppress

# Xóa file cũ nếu có, nếu không có cũng không báo lỗi
with suppress(FileNotFoundError):
    os.remove("cache_temp.tmp")

2. contextlib.closing: Tự động đóng các object đời cũ

Một số object cũ có phương thức close() nhưng lại không hỗ trợ từ khóa with. Bạn có thể dùng closing để “ép” chúng tuân thủ giao thức Context Manager.

from contextlib import closing
from urllib.request import urlopen

with closing(urlopen('https://www.google.com')) as page:
    content = page.read()
    print(f"Đã tải {len(content)} bytes")
# Kết nối tự đóng ngay khi thoát khối with

3. contextlib.ExitStack: Quản lý hàng loạt tài nguyên

Đây là tính năng mình tâm đắc nhất. Giả sử bạn cần mở cùng lúc 10 file log để gộp dữ liệu. Việc lồng 10 khối with sẽ tạo ra một “kim tự tháp” code cực kỳ khó nhìn. ExitStack giúp bạn quản lý danh sách động các tài nguyên một cách phẳng và gọn.

from contextlib import ExitStack

def merge_logs(log_files, output_path):
    with ExitStack() as stack:
        # Mở toàn bộ file trong danh sách cùng lúc
        files = [stack.enter_context(open(fname)) for fname in log_files]
        with open(output_path, 'w') as out:
            for f in files:
                out.write(f.read())
# Tất cả file được đóng gọn gàng khi kết thúc hàm

Xử lý lỗi như một chuyên gia

Khi viết Context Manager, điều quan trọng là không để script bị “treo” khi có ngoại lệ (exception). Nếu có lỗi xảy ra bên trong khối with, nó sẽ được ném ra ngay tại vị trí yield trong generator của bạn.

Bạn có thể bắt lỗi này để rollback database hoặc log lại thông tin trước khi để nó tiếp tục lan truyền:

@contextmanager
def db_transaction(conn):
    print("Bắt đầu Transaction...")
    try:
        yield conn
        conn.commit()
    except Exception as e:
        conn.rollback()
        print(f"Lỗi: {e}. Đã rollback dữ liệu.")
        raise 
    finally:
        conn.close()

Một mẹo nhỏ là hãy kết hợp với module logging thay vì dùng print. Điều này giúp bạn theo dõi xem có tài nguyên nào đang bị chiếm dụng quá lâu trong môi trường production hay không.

Hãy nhớ quy tắc vàng: Mở ở đâu, dọn dẹp ở đó. Đừng bao giờ chủ quan cho rằng code bên trong with sẽ luôn chạy mượt mà. Luôn bọc yield trong try...finally là cách tốt nhất để bảo vệ hệ thống của bạn.

Tóm lại, contextlib không chỉ giúp code đẹp hơn mà còn là lá chắn giúp ứng dụng của bạn ổn định hơn. Nếu bạn đang xây dựng các hệ thống chạy 24/7, hãy áp dụng contextlib ngay hôm nay để tránh những lỗi rò rỉ tài nguyên đáng tiếc.