Chuyển đến nội dung chính

Viết Context Manager tùy chỉnh trong Python

Tìm hiểu các khía cạnh nâng cao của quản lý tài nguyên trong Python bằng cách nắm vững cách viết context manager tùy chỉnh.

Khám phá cùng AI

Động lực

Có thể bạn đã quen với câu lệnh with, một cách ngắn gọn để xử lý tài nguyên trong Python. Nhưng bạn đã bao giờ tự hỏi nó hoạt động như thế nào ở bên trong chưa? Sức mạnh của từ khóa with đến từ context manager.

Context manager là một mẫu thiết kế nền tảng trong Python, cung cấp cách tiếp cận có cấu trúc cho việc quản lý tài nguyên. Chúng đảm bảo tài nguyên được cấp phát, sử dụng đúng cách, và cuối cùng được giải phóng hoặc dọn dẹp, ngay cả khi có lỗi hoặc ngoại lệ.

Điều này đặc biệt quan trọng khi làm việc với các tài nguyên như tệp, kết nối mạng, hoặc handle cơ sở dữ liệu. Bằng cách dùng context manager, chúng ta có thể viết mã gọn gàng và đáng tin cậy hơn, không còn lo lắng chuyện quên đóng tệp hoặc nhả khóa.

Trong hướng dẫn này, chúng ta sẽ đi xa hơn các context manager mặc định của Python và học cách viết của riêng mình.

Tìm hiểu Context Manager trong Python

Bên dưới, context manager là các đối tượng định nghĩa hai phương thức đặc biệt: __enter__ và __exit__. Phương thức __enter__ được gọi khi bạn đi vào khối with, và giá trị trả về của nó được gán cho một biến trong khối đó. Ngược lại, phương thức __exit__ được gọi khi khối with kết thúc, bất kể kết thúc bình thường hay có ngoại lệ.

Cấu trúc này đảm bảo việc xử lý tài nguyên đúng đắn. Hãy xem một số context manager dựng sẵn trong Python để hiểu rõ hơn:

1. Quản lý tệp

Xem ví dụ kinh điển về việc mở một tệp:

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

Ở đây, open('file.txt', 'r') đóng vai trò là context manager. Khi bạn vào khối with, phương thức enter của đối tượng tệp được gọi, mở tệp và gán nó cho biến file. Sau đó bạn có thể dùng file.read() để truy cập nội dung tệp.

Tiếp theo, phương thức __exit__ của đối tượng tệp chắc chắn sẽ được gọi khi khối with kết thúc, kể cả khi có ngoại lệ. Phương thức này đảm nhiệm việc đóng tệp, đảm bảo bạn không để sót các handle tệp đang mở.

2. Khóa luồng

Vượt ra ngoài tệp, context manager cũng có thể được dùng để đồng bộ hóa luồng bằng threading.Lock():

import threading

lock = threading.Lock()
with lock:
   # Critical section
   print("This code is executed under lock protection.")

Ở đây, lock là một đối tượng threading.Lock, cũng là một context manager. Khi vào khối with, phương thức __enter__ sẽ chiếm khóa, đảm bảo chỉ một luồng có thể thực thi vùng tới hạn tại một thời điểm. Cuối cùng, phương thức __exit__ sẽ nhả khóa khi thoát khỏi khối with, cho phép các luồng khác tiếp tục.

3. Kết nối cơ sở dữ liệu

Tương tự, context manager có thể quản lý kết nối cơ sở dữ liệu:

import sqlite3

with sqlite3.connect('database.db') as connection:
    cursor = connection.cursor()
    cursor.execute("SELECT * FROM table")
    rows = cursor.fetchall()

Lời gọi sqlite3.connect('database.db') là một context manager. Khi vào khối with, một kết nối tới cơ sở dữ liệu được thiết lập và gán cho connection. Sau đó bạn có thể dùng một cursor để tương tác với cơ sở dữ liệu. Phương thức __exit__ đảm bảo kết nối được đóng khi khối with kết thúc, ngăn ngừa rò rỉ tài nguyên.

4. Ổ cắm mạng

Context manager thậm chí có thể xử lý giao tiếp mạng:

import socket

with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
    s.connect(('localhost', 8080))
    s.sendall(b'Hello, world')
    data = s.recv(1024)

Ở đây, socket.socket(socket.AF_INET, socket.SOCK_STREAM) tạo một đối tượng socket đóng vai trò context manager. Phương thức __enter__ tạo socket, và trong khối with, bạn có thể kết nối, gửi và nhận dữ liệu. Phương thức __exit__ đảm bảo socket được đóng đúng cách khi xong việc.

5. Quét thư mục

Lần này, os.scandir('.') cung cấp cách lặp qua các mục trong thư mục:

with os.scandir('.') as entries:
   for entry in entries:
       print(entry.name)

os.scandir('.') đóng vai trò context manager ở đây. Phương thức __enter__ mở một phiên quét thư mục, và bạn có thể lặp qua các mục trong khối with. Phương thức __exit__ dọn dẹp phiên quét khi thoát.

Như bạn thấy, câu lệnh with, được hỗ trợ bởi context manager, đơn giản hóa quản lý tài nguyên bằng cách tự động cấp phát và thu hồi.

Viết context manager tùy chỉnh trong Python

Giờ hãy xem cách viết các manager tùy chỉnh của riêng bạn để kiểm soát chi tiết việc quản lý tài nguyên.

Có hai cách chính để viết context manager tùy chỉnh: dựa trên lớp (class-based) và dựa trên hàm (function-based).

Cách tiếp cận dựa trên lớp

Cách dựa trên lớp là phương pháp có cấu trúc và linh hoạt nhất để viết context manager. Bạn định nghĩa một lớp triển khai các phương thức đặc biệt __enter__ và __exit__.

Hãy xem ví dụ về lớp Timer đo thời gian thực thi:

import time


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

   def __exit__(self, exc_type, exc_value, traceback):
       self.end_time = time.time()
       elapsed_time = self.end_time - self.start_time
       print(f"Elapsed time: {elapsed_time} seconds")


# Example usage
if __name__ == "__main__":
   with Timer() as timer:
       # Code block to measure the execution time
       time.sleep(2)  # Simulate some time-consuming operation
Elapsed time: 2.002082347869873 seconds

Ở đây, lớp Timer định nghĩa phương thức __enter__ để ghi nhận thời điểm bắt đầu khi bạn vào khối with. Nó trả về self để cho phép truy cập đối tượng trong khối. Phương thức __exit__ tính thời gian đã trôi qua khi thoát khỏi khối with và in ra.

Cách tiếp cận dựa trên hàm

Nếu bạn thích cách ngắn gọn hơn, có thể đánh đổi một chút tính linh hoạt, thì cách dựa trên hàm là lựa chọn tốt. Bạn dùng decorator contextmanager từ mô-đun contextlib để biến bất kỳ hàm nào thành một context manager.

Hãy xem cách hoạt động:

import time
from contextlib import contextmanager


@contextmanager
def timer():
   start_time = time.time()
   yield
   end_time = time.time()
   elapsed_time = end_time - start_time
   print(f"Elapsed time: {elapsed_time} seconds")


# Example usage
if __name__ == "__main__":
   with timer():
       time.sleep(2)
Elapsed time: 2.0020740032196045 seconds

Decorator @contextmanager biến hàm timer thành một context manager. Bên trong hàm, start_time được ghi nhận, và câu lệnh yield tạm dừng thực thi, cho phép mã trong khối with chạy.

Cuối cùng, chức năng tương đương __exit__ được thực hiện bằng cách lấy thời điểm kết thúc và in thời gian đã trôi qua.

Về bản chất, bạn viết logic cho __enter__ trước từ khóa yield còn logic cho __exit__ nằm sau đó. Cả hai cách đều đạt cùng kết quả, lựa chọn tùy thuộc sở thích của bạn về cấu trúc và khả năng đọc.

Ví dụ thực tiễn về Context Manager

Giờ khi bạn đã hiểu cách viết context manager tùy chỉnh trong Python, hãy xem một số ví dụ thực tế nơi chúng cực kỳ hữu ích.

Quản lý thao tác với tệp

Thao tác với tệp là nhiệm vụ phổ biến trong nhiều ứng dụng. Dù đọc hay ghi tệp, quản lý đúng cách đảm bảo tính toàn vẹn dữ liệu và hiệu quả tài nguyên. Hãy xét một tình huống muốn ghi văn bản vào tệp bằng context manager tùy chỉnh.

class FileManager:
   def __init__(self, filename, mode):
       self.filename = filename
       self.mode = mode
       self.file = None

   def __enter__(self):
       self.file = open(self.filename, self.mode)
       return self.file

   def __exit__(self, exc_type, exc_value, traceback):
       if self.file:
           self.file.close()


# Example usage
if __name__ == "__main__":
   with FileManager("example.txt", "w") as file:
       file.write("Hello, world!\n")

Trong ví dụ này, chúng ta định nghĩa lớp FileManager để bắt chước hàm open trong Python nhằm xử lý thao tác tệp. Hãy phân tích từng phương thức:

  • __init__(self, filename, mode): Phương thức khởi tạo của lớp. Nó khởi tạo đối tượng FileManager với tên tệp và chế độ được cung cấp, chỉ định tệp sẽ mở và chế độ mở (ví dụ: đọc, ghi, nối).
  • __enter__(self): Phương thức được gọi khi vào khối with như thường lệ. Nó mở tệp theo filename ở chế độ mode. Đối tượng tệp mở được gán cho self.file và trả về để dùng trong khối with.
  • __exit__(self, exc_type, exc_value, traceback): Phương thức được gọi khi thoát khối with, bất kể có ngoại lệ trong khối hay không. Nó đảm bảo tệp được đóng đúng cách bằng cách gọi self.file.close() nếu self.file không phải None. Điều này ngăn rò rỉ tài nguyên và đảm bảo dọn dẹp.

Nếu bạn kiểm tra nội dung của example.txt sau khi chạy đoạn mã trên, bạn sẽ thấy văn bản đã được ghi:

$ cat example.txt
Hello, world!

Tất nhiên, ưu điểm của lớp FileManager là bạn có thể triển khai bất kỳ logic tùy chỉnh nào để làm việc với tệp. Chẳng hạn, bạn có thể chỉnh sửa để lớp lưu log những tệp nào được mở hoặc đóng vào một tệp văn bản.

import datetime


class FileManager:
   def __init__(self, filename, mode, log_filename="file_log.txt"):
       self.filename = filename
       self.mode = mode
       self.log_filename = log_filename
       self.file = None

   def log_action(self, action):
       timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
       with open(self.log_filename, "a") as log_file:
           log_file.write(f"{timestamp} - {action}: {self.filename}\n")

   def __enter__(self):
       self.file = open(self.filename, self.mode)
       self.log_action("Opened")
       return self.file

   def __exit__(self, exc_type, exc_value, traceback):
       if self.file:
           self.file.close()
           self.log_action("Closed")


# Example Usage
with FileManager("example.txt", "r") as file:
   content = file.read()
   print(content)

Hello, world!

Trong phiên bản đã chỉnh sửa này:

  • Mô-đun datetime được nhập để làm việc với dấu thời gian.
  • Thêm phương thức log_action để ghi log thao tác với tệp.
  • Thêm tham số log_filename để chỉ định tên tệp log. Mặc định là 'file_log.txt'.
  • Tên tệp và hành động (mở hoặc đóng) được nối thêm vào tệp log bằng phương thức log_action trong các phương thức __enter__ và __exit__.
  • Bên trong phương thức log_action, datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") được dùng để lấy dấu thời gian hiện tại theo định dạng "YYYY-MM-DD HH:MM:SS".
  • Dấu thời gian sau đó được đưa vào log cùng với tên tệp và hành động.

Giờ, nếu chúng ta in nội dung của file_log.txt, sẽ thấy các dòng log:

$ cat file_log.txt
2024-04-03 12:37:02 - Opened: example.txt
2024-04-03 12:37:02 - Closed: example.txt

Một trường hợp sử dụng thực tế khác có thể triển khai như một context manager tùy chỉnh là công cụ quản lý cơ sở dữ liệu SQLite:

import sqlite3


class DatabaseConnection:
   def __init__(self, database_name):
       self.database_name = database_name
       self.connection = None

   def __enter__(self):
       self.connection = sqlite3.connect(self.database_name)
       return self.connection

   def __exit__(self, exc_type, exc_value, traceback):
       if self.connection:
           self.connection.close()

Một lần nữa, điều này bắt chước hàm sqlite3.connect, nhưng có thể cải tiến thêm bằng cách bổ sung logic quản lý cơ sở dữ liệu tùy chỉnh. Hãy kiểm tra manager hoạt động bằng vài hàm mẫu:

# Example usage
def create_table():
   with DatabaseConnection("example.db") as connection:
       cursor = connection.cursor()
       cursor.execute(
           """CREATE TABLE IF NOT EXISTS users (
                           id INTEGER PRIMARY KEY,
                           username TEXT,
                           email TEXT)"""
       )


def insert_data(username, email):
   with DatabaseConnection("example.db") as connection:
       cursor = connection.cursor()
       cursor.execute(
           "INSERT INTO users (username, email) VALUES (?, ?)", (username, email)
       )
       connection.commit()


def fetch_data():
   with DatabaseConnection("example.db") as connection:
       cursor = connection.cursor()
       cursor.execute("SELECT * FROM users")
       return cursor.fetchall()


if __name__ == "__main__":
   create_table()
   insert_data("john_doe", "[email protected]")
   insert_data("jane_doe", "[email protected]")
   users = fetch_data()
   print("Users in the database:")
   for user in users:
       print(user)

Users in the database:
(1, 'john_doe', '[email protected]')
(2, 'jane_doe', '[email protected]')
(3, 'john_doe', '[email protected]')
(4, 'jane_doe', '[email protected]')

Chúng ta có thể viết mã dễ bảo trì hơn nhiều bằng cách tách các thao tác cơ sở dữ liệu thành các hàm và dùng context manager tùy chỉnh vừa tạo.

Chủ đề nâng cao liên quan đến Context Manager trong Python

Trong phần này, chúng ta sẽ bàn về các chủ đề nâng cao liên quan đến context manager trong Python, gồm xử lý lỗi và lồng ghép context manager.

Xử lý lỗi

Xử lý lỗi là một khía cạnh thiết yếu của lập trình, và context manager trong Python cung cấp cách tiện lợi để quản lý lỗi trong các thao tác quản lý tài nguyên.

Trong ví dụ lớp FileManager, xử lý lỗi có thể được triển khai trong phương thức __exit__ để đảm bảo dọn dẹp đúng cách ngay cả khi có ngoại lệ. Ví dụ, nếu xảy ra lỗi khi ghi dữ liệu vào tệp (khi đang ở trong khối with), chúng ta vẫn muốn đảm bảo tệp được đóng để tránh rò rỉ tài nguyên. Ta có thể sửa phương thức __exit__ như sau:

def __exit__(self, exc_type, exc_value, traceback):
   if self.file:
       self.file.close()
       if exc_type is not None:
           print(f"An error occurred: {exc_value}")

Ở đây, trước tiên chúng ta đóng tệp bằng self.file.close() để đảm bảo dọn dẹp. Sau đó, kiểm tra xem có ngoại lệ xảy ra không (exc_type is not None). Nếu có, chúng ta in thông báo lỗi để cung cấp thông tin về sự cố.

Cụ thể:

  • exc_type là lớp ngoại lệ, như AttributeError và IndexError.
  • exc_value là thể hiện ngoại lệ kèm thông điệp.
  • traceback là toàn bộ vết ngăn xếp của lỗi.

Hãy chạy lớp FileManager với một đoạn mã cố ý sai:

with FileManager("example.txt", "r") as file:
   # Simulate a sample error
   content = file.read()
   print(content[150])

An error occurred: string index out of range

Vì vậy, để triển khai logic tùy chỉnh khi xử lý ngoại lệ trong các manager của bạn, đây là một quy trình mẫu bạn có thể làm theo:

def __exit__(self, exc_type, exc_value, traceback):
   # Perform cleanup
   ...

   # Handle the exception
   if isinstance(exc_value, AttributeError):
       # Handle AttributeError here
       print(f"Exception message: {exc_value}")
   # Handle another exception
   elif isinstance(exc_value, AnotherError):
       ...

Lồng ghép context manager

Lồng ghép context manager cho phép bạn quản lý đồng thời nhiều tài nguyên, mỗi cái có context manager riêng, trong một câu lệnh with duy nhất. Điều này đặc biệt hữu ích khi xử lý các thao tác phức tạp liên quan đến nhiều tài nguyên. Hãy xem cách chúng ta có thể lồng ghép context manager trong các ví dụ của mình.

Giả sử bạn muốn mở một tệp và thực hiện các thao tác cơ sở dữ liệu trong cùng một ngữ cảnh. Có thể, bạn muốn đưa nội dung của một tệp CSV vào cơ sở dữ liệu dưới dạng các hàng. Chúng ta có thể đạt được điều này bằng cách lồng ghép context manager FileManager và DatabaseConnection trong một câu lệnh with duy nhất:

if __name__ == "__main__":
   with FileManager("example.txt", "w") as file, DatabaseConnection("example.db") as connection:
       ...

Trong ví dụ này, context manager FileManager và DatabaseConnection được lồng trong cùng một câu lệnh with. Điều này cho phép chúng ta mở tệp và thiết lập kết nối cơ sở dữ liệu đồng thời. Các thao tác ghi tệp và cơ sở dữ liệu tiếp theo được thực hiện trong ngữ cảnh kết hợp này.

Một điều cần lưu ý là tính dễ đọc của mã. Tôi khuyến nghị số lượng ngữ cảnh đồng thời không vượt quá hai.

Kết luận

Context manager trong Python cung cấp cách tiếp cận có cấu trúc để quản lý tài nguyên trong câu lệnh with, đảm bảo cấp phát và thu hồi tài nguyên đúng cách.

Xuyên suốt hướng dẫn này, chúng ta đã học những điều cơ bản về context manager, đi qua các ví dụ thực tiễn, và bàn về các chủ đề nâng cao như xử lý lỗi và lồng ghép. Bằng cách thành thạo viết context manager tùy chỉnh, bạn có thể viết mã gọn gàng, tin cậy hơn, cải thiện xử lý lỗi và quản lý tài nguyên hiệu quả.

Chủ đề viết context manager tùy chỉnh nằm trong danh mục “Python nâng cao”. Dưới đây là một số tài nguyên liên quan để nâng cao kỹ năng của bạn:

Bexruz (Bex) Tuychiev

Tác giả

Bexruz (Bex) Tuychiev

LinkedIn

Tôi là người sáng tạo nội dung về khoa học dữ liệu với hơn 2 năm kinh nghiệm và là một trong những tài khoản có lượng người theo dõi lớn nhất trên Medium. Tôi thích viết các bài chuyên sâu về AI và ML với chút giọng điệu mỉa mai, vì bạn cũng phải làm gì đó để chúng bớt khô khan. Tôi đã xuất bản hơn 130 bài viết và một khóa học trên DataCamp, cùng với một khóa nữa đang được thực hiện. Nội dung của tôi đã được hơn 5 triệu lượt xem, trong đó 20 nghìn người đã trở thành người theo dõi trên cả Medium và LinkedIn. 

Chủ đề

Python

Liên quan

Xem thêm →