Python context managers

Trong Python, context manager là một phương thức cho phép bạn cấp phát và sử dụng tài nguyên một cách hiệu quả. Context manager được sử dụng rộng rãi thông qua câu lệnh with. Ví dụ:

with open('foo', 'w') as f:
    f.write('Hora! We opened this file')

Đoạn code trên mở một file, ghi dữ liệu và đóng file lại. Nếu có bất kỳ lỗi gì xảy ra, thì file cũng luôn luôn được đảm bảo là đã đóng. Đoạn code trên nếu viết mà không sử dụng context manager thì sẽ trông như dưới đây:

f = open('foo', 'w')
try:
    f.write('Hora! We opened this file')
finally:
    f.close()

So sánh hai cách viết này thì chúng ta đã thấy rất rõ ràng rằng, context manager cho chúng ta cách viết code ngắn gọn hơn hẳn. Lệnh with cho chúng ta bảo đảm rằng file luôn luôn được đóng mà không cần biết những logic xử lý bên trong.

Hẳn là các bạn đã rất quen thuộc với những đoạn code trên, đặc biệt khi bạn đã từng nghe nói đến “Idiomatic Python”. Nhưng liệu bạn đã chắc chắn hiểu được cách làm việc chính xác với file và lý do tại sao đó lại là cách đúng không? Hoặc đơn giản hơn, bạn có biết khi nào thì mình đã thao tác sai không.

Nếu câu trả lời là không, thì bài viết này chính là dành cho bạn.

Context manager thường được sử dụng để lock các tài nguyên (trường hợp mở và đóng file là một ví dụ kinh điển cho việc này).

Quản lý tài nguyên

Tính năng quan trọng nhất và cũng là phổ biến nhất của context manager là để quản lý tài nguyên một cách chính xác. Quay lại với việc đọc và ghi file ở ví dụ trên, tại sao chúng ta phải sử dụng context manager. Mỗi khi mở một file để đọc hoặc ghi, một tài nguyên của hệ thống, trong trường hợp này là file descriptor sẽ đã bị tiêu tốn để chúng ta có thể thao tác. Thật không may là tài nguyên này lại là hữu hạn. Mỗi hệ điều hành đều có giới hạn nhất định cho số lượng file có thể mở cùng một lúc.

Không tin ư, bạn hãy xem ví dụ sau:

>>> files = []
>>> for _ in range(10000):
...     files.append(open('foo', 'w'))
...
Traceback (most recent call last):
  File "<stdin>", line 2, in <module>
OSError: [Errno 24] Too many open files: 'foo'

Ngoài lề một chút, file descriptor thực chất là một số nguyên. Khi bạn mở một file, hệ điều hành sẽ tạo ra một entry (có thể ở kernel) lưu trữ những thông tin liên quan đến file được mở. Mỗi entry sẽ được gán với 1 số nguyên (và số này sẽ là duy nhất), cho phép người dùng thông qua đó để thao tác với file.

Thực chất người dùng đang thao tác một cách gián tiếp thông qua file descriptor (và thông qua entry của kernel) với các dữ liệu thật sự được ghi ở bộ nhớ ngoài. Việc này mang lại nhiều lợi ích như có thể chia sẻ file cho nhiều tiến trình khác nhau cũng như duy trì bảo mật cho các file đó.

Thực ra, hằng ngày bạn đều đang làm việc với file descriptor mà có thể bạn cũng không nhận ra. Hệ điều hành đã gán sẵn một số file descriptor như 0 cho bàn phím, 1 cho màn hình, v.v… Và mọi thao tác chúng ta làm với máy tính đều thông qua các file descriptor này. Bạn nghĩ sao về việc chuyển hướng của câu lệnh Linux như thế này:

$ sort < file_list > sorted_file_list 2>&1

Tương tự như vậy, khi bạn mở một socket, một socket descriptor cũng sẽ được sử dụng.

Quay trở lại với nội dung của bài viết. Vậy điều gì xảy khi code Python của bạn mở file mà không đóng nó lại. Rất hiển nhiên, chúng ta đã mất một file descriptor mà chúng ta sẽ không bao giờ cần đến nó nữa. Điều này đồng nghĩa với việc, số file chúng ta có thể thao tác sẽ ít dần đi, do số lượng của chúng bị giới hạn. Mà việc “quên” không đóng file nhiều khi xảy ra khá thường xuyên, và tích tiểu thành đại, đến một lúc nào đó bạn không thể mở thêm file nào nữa.

Bạn có thể dùng lệnh ulimit -n để kiểm tra xem hệ thống của mình cho phép mở tối đa bao nhiêu file cùng lúc.

Tất nhiên là mọi vấn đề đều có thể giải quyết được. Vẫn với ví dụ trên, chúng ta có thể xử lý bằng cách đóng từng file một như sau:

>>> files = []
>>> for _ in range(10000):
...     f = open('foo', 'w')
...     f.close()
...     files.append(f)
...
>>>

Quản lý tài nguyên hiệu quả hơn

Trên đây là một cách giải quyết tuy vẫn hoạt động tốt nhưng không được thông minh cho lắm. Trong những hệ thống phức tạp hơn, rất khó để đảm bảo rằng tất cả các file đã được đóng lại khi không dùng đến nữa.

Giả sử trong quá trình thao tác, chúng ta gặp phải một exception nào đó thì phải làm thế nào đây. Bắt exception và xử lý riêng sao? Phải bắt những exception nào mới thì gọi là đủ?

Hoặc kể cả không có exception nhưng hàm đã return trước khi file kịp close thì sao? Trong những trường hợp phức tạp như vậy, làm thế nào để chúng ta “nhớ” phải đóng file lại. Câu trả lời là khó tới gần như không thể (thật phũ phàng).

Trong nhiều ngôn ngữ, lập trình viên cần phải sử dụng cấu trúc kiểu như try ... except ... finally ... để đảm bảo rằng file sẽ được đóng. Rất may mắn, Python đã nghĩ đến những khó khăn này của chúng ta và đưa cho chúng ta một phương thức dễ dàng để làm những việc đó - context manager.

Nói một cách ngắn gọn, chúng ta cần một phương thức càng đơn giản càng tốt để đảm bảo các tài nguyên được dọn dẹp cẩn thận dù có xảy ra bất cứ chuyện gì đi chăng nữa. Và context manager sẽ cung cấp cho chúng ta tính năng này:

with something_that_returns_a_context_manager() as my_resource:
    do_something(my_resource)
    ...
    print('done using my_resource')

Đơn giản vậy đó. Bằng cách sử dụng with, chúng ta sẽ đưa mọi thứ vào trong một context manager. Chúng ta gán context manager này cho một biến, và biến đó chỉ tồn tại khi block sau đó được thực thi. Điều này giống như chúng ta tạo một hàm, nó sẽ gọi một số thao tác và khi kết thúc, nó sẽ tự dọn dẹp những gì nó tạo ra.

Một số context manager hữu ích khác

Context manager thực sự rất cần thiết trong Python, và nó đã có mặt trong thư việc chuẩn. Một số context manager có thể bạn đã từng làm việc là zipfile.ZipFiles, subprocess.Popen, tarfile.TarFile, telnetlib.Telnet, pathlib.Path, v.v… Thậm chí, Lock của threading cũng là context manager. Trên thực tế, tất cả những tài nguyên mà chúng ta cần close sau khi sử dụng đều (và rất nên) là context manager.

Việc sử dụng Lock tương đối đặc biệt một chút. Trong trường hợp này, tài nguyên là một mutex. Sử dụng context manager sẽ phòng tránh được deadlock trong lập trình multithread nếu chúng ta sử dụng khóa mà không bao giờ mở nó. Hãy xem xét ví dụ sau:

>>> from threading import Lock
>>> lock = Lock()
>>> def do_something_dangerous():
...     lock.acquire()
...     raise Exception('OOPS! I forgot this code could raise exceptions')
...     lock.release()
...
>>> try:
...     do_something_dangerous()
... except:
...     print('Got an exception')
...
Got an exception
>>> lock.acquire()

Với code trên, rõ ràng là lock.release() sẽ không bao giờ được gọi, và do đó, mọi tiến trình sẽ gặp deadlock và chết cứng ở đó (lock.acquire() sẽ không bao giờ kết thúc). Rất may mắn, với context manager, điều này có thể sửa chữa được:

>>> from threading import Lock
>>> lock = Lock()
>>> def do_something_dangerous():
...     with lock:
...         raise Exception('oops I forgot this code could raise exceptions')
...
>>> try:
...     do_something_dangerous()
... except:
...     print('Got an exception')
...
Got an exception
>>> lock.acquire()
True
>>> print('We can get here')
We can get here
>>>

Trên thực tế, không có cách nào để gây ra deadlock nếu sử dụng context manager. Và đây là điều chúng ta đang cần.

Trong phần tiếp theo, chúng ta sẽ tìm hiểu cách cài đặt một context manager, và qua đó, chúng ta sẽ hiểu hơn về cách thức một context manager hoạt động.

Cài đặt context manager như một class

Có nhiều cách khác nhau để cài đặt một context manager. Cách đơn giản nhất là cài đặt một class với hai phương thức vô cùng đặc biệt: __enter__ và __exit__. Phương thức __enter__ sẽ trả về tài nguyên cần quản lý (ví dụ như file đang được mở) và __exit__ sẽ làm việc dọn dẹp hệ thống.

Hãy xem xét ví dụ sau về một context manager khi làm việc với file:

>>> class File:
...     def __init__(self, file_name, method):
...         self.file_obj = open(file_name, method)
...     def __enter__(self):
...         return self.file_obj
...     def __exit__(self, type, value, traceback):
...         self.file_obj.close()
...
>>>

Class trên cũng như nhiều class khác, phương thức __init__ để để khởi tạo đối tượng, trong trường hợp này là khởi tạo tên file cần mở cùng với mode (đọc/ghi) của nó. Phương thức __enter__ mở file và trả về đối tượng file để thao tác với file đó trong khi __exit__ chỉ đơn giản là đóng file lại.

Với hai phương thức __enter__ và __exit__, chúng ta có thể sử dụng class này cùng với with:

>>> with File('foo', 'w') as f:
...     f.write('Hora! We opened this file')
...
25

Phương thức __exit__ bắt buộc phải có 3 tham số. Dưới đây là những gì thực sự xảy ra khi chúng ta gọi context manager:

• Câu lệnh with lưu phương thức __exit__ của class File
• Câu lệnh này gọi phương thức __enter__ của class File
• Phương thức __enter__ mở file và trả về object để thao tác với file đó
• Object được trả về được truyền cho biến f
• Chúng ta thao tác với file bằng cách ghi dữ liệu f.write
• Khi kết thúc block, câu lệnh with gọi phương thức __exit__
• Phương thức __exit__ đóng file cho chúng ta

Xử lý exception

Trong cài đặt đơn giản trên, chúng ta đã bỏ qua 3 tham số type, value, traceback của phương thức __exit__. Tuy nhiên, trong quá trình thực thi block lệnh ở trên, nếu xảy ra một exception, Python sẽ chuyển những thông tin type, value và traceback của exception này tới phương thức __exit__. Điều đó giúp chúng ta có thể tùy biến phương thức __exit__ để xử lý những vấn đề có thể xảy ra trong quá trình thực thi. Trong trường hợp của chúng ta, chúng ta chỉ cần đóng file và không cần quan tâm đến exception này.

Nhưng chuyện gì sẽ xảy ra nếu bản thân file object gặp phải một exception? Ví dụ, khi chúng ta thử gọi một phương thức không tồn tại:

>>> with File('foo', 'w') as f:
...     f.undefined_function('Oops! I called an unknown method')
...
Traceback (most recent call last):
  File "<stdin>", line 2, in <module>
AttributeError: '_io.TextIOWrapper' object has no attribute 'undefined_function'

Dưới đây là quy trình những gì đã xảy ra khi có lỗi xảy ra:

• type, value, traceback của lỗi đó được truyền cho __exit__
• Trong phương thức __exit__, chúng ta có thể tùy ý xử lý exception đó
• Nếu __exit__ trả về True thì exception đã được xử lý hoàn toàn.
• Nếu không, exception sẽ tiếp tục được raise bởi lệnh with

Trong trường hợp của chúng ta, phương thức __exit__ không trả về bất cứ thứ gì, do đó, lệnh with sẽ raise exception.

Chúng ta có thể tạm xử lý exception như sau:

>>> class File:
...     def __init__(self, file_name, method):
...         self.file_obj = open(file_name, method)
...     def __enter__(self):
...         return self.file_obj
...     def __exit__(self, type, value, traceback):
...         print("Exception has been handled")
...         self.file_obj.close()
...         return True
...
>>> with File('foo', 'w') as f:
...     f.undefined_function()
...
Exception has been handled
>>>

Phương thức __exit__ trả về True, do đó, không có exception nào được raise bởi lệnh with.

Có nhiều cách để cài đặt context manager. Trên đây là một cách đơn giản và dễ hiểu nhất. Trong phần tiếp theo, chúng ta sẽ tìm hiểu thêm một số phương pháp cài đặt nữa.

Sử dụng contextlib cài đặt context manager

Context manager quả là tiện lợi và hữu ích vô cùng. Do đó, trong thư viện chuẩn của Python có hẳn module contextlib với rất nhiều công cụ để tạo và làm việc với context manager.

Chúng ta có thể cài đặt context manager bằng cách sử dụng decorator và generator. contextlib cung cấp cho chúng ta decorator @contextmanager để decorate các hàm generator chỉ gọi yield đúng một lần duy nhất. Với decorator này, tất cả những gì diễn ra trước yield đều được coi là thao tác của phương thức __enter__. Những gì dễn ra sau đó được coi là của phương thức __exit__.

Hãy xem xét ví dụ của chúng ta về quản lý file khi dùng contextlib:

>>> from contextlib import contextmanager
>>> @contextmanager
... def open_file(path, mode):
...     f = open(path, mode)
...     yield f
...     f.close()
...
>>> files = []
>>> for _ in range(10000):
...     with open_file('foo', 'w') as f:
...         files.append(f)
...
>>> for f in files:
...     if not f.closed:
...         print('not closed')
...
>>>

Như chúng ta đã thấy, việc cài đặt context manager đã ngắn gọn hơn rất nhiều. Chúng ta chỉ cần mở file, yield đối tượng đó và đóng nó lại. Mọi việc còn lại sẽ do decorator @contextmanager đảm nhiệm.

Và ví dụ thực tế cho thấy rằng các file của chúng ta đã được quản lý tốt, tất cả chúng đã được đóng lại đầy đủ. Tuy nhiên, cách thức cài đặt tiện lợi này yêu cầu chúng ta phải có chút hiểu biết về decorator, generator cũng như lệnh yield. Có thể tóm tắt quá trình tạo context manager trên như sau:

• Python tìm thấy yield, hàm này là một generator chứ không phải hàm thông thường.
• Với decorator @contextmanager, hàm open_file sẽ được truyền là tham số cho hàm contextmanager.
• Hàm contextmanager trả về generator được bọc trong object của GeneratorContextManager.
• Object GeneratorContextManager được gán cho hàm open_file. Do đó, khi chúng ta gọi hàm này, thực ra chúng ta đang làm việc với object GeneratorContextManager.

Python docs còn một ví dụ khác thú vị hơn:

>>> from contextlib import contextmanager
>>> @contextmanager
... def tag(name):
...     print("<%s>" % name)
...     yield
...     print("</%s>" % name)
...
>>> with tag('h1'):
...     print('foo')
...
<h1>
foo
</h1>

Trong tất cả các trường hợp trên, chúng ta cũng không hề xử lý exception, vì vậy, context manager của chúng ta sẽ hoạt động giống như code đầu tiên.

Một công cụ tiện lợi khác của contextlib là ContextDecorator. Nó cho phép chúng ta cài đặt các context manager theo kiểu class. Nhưng với việc kế thừa từ class ContextDecorator, bạn có thể sử dụng context manager với lệnh with thông thường, hoặc sử dụng nó như một decorator dùng để decorate các hàm khác. Chúng ta có thể xem xét ví dụ sau (tương tự như ví dụ tag HTML ở trên):

>>> from contextlib import ContextDecorator
>>> class tag(ContextDecorator):
...     def __init__(self, name):
...         self.name = name
...     def __enter__(self):
...         print('<%s>' % self.name)
...         return self
...     def __exit__(self, *exc):
...         print('</%s>' % self.name)
...         return False
...
>>> with tag('h1'):
...     print('this is not html')
...
<h1>
this is not html
</h1>
>>> @tag('h1')
... def content():
...     print('this is another non-html content')
...
>>> content()
<h1>
this is another non-html content
</h1>
>>>

Kết luận

Bài viết trình bày những hiểu biết của tôi về context manager của Python, cách nó hoạt động và hỗ trợ cho chúng ta trong công việc lập trình. Như các bạn đã thấy, chúng ta có thể làm rất nhiều thứ với context manager. Mục đích cao nhất của nó thì không bao giờ thay đổi: quản lý hiệu quả các tài nguyên.

Chúng ta không chỉ có thể dùng context manager mà còn có thể tự cài đặt context manager cho riêng mình. Hãy sử dụng context manager và làm cho cuộc sống dễ chịu hơn.