Context managers
Context managers are a powerful feature for managing resources such as files, database connections, and network sockets.
A problem that can be solved with a context manager arises if, while reading a file, the program crashes and doesn't reach the close()
method. This could lead to file corruption. Of course, we could solve it with try
, except
, and finally
, but a better way would be to use the with
keyword. It uses a context manager under the hood, which means it already has the "try
, except
, finally
" structure implemented (a file will be properly closed even if when reading from it, the program crashes). We can customize this context manager with the __enter__()
and __exit__()
dunder methods if we want to fully control what happens while entering the file and during its closeup.
class File:
def __init__(self, file_name, method):
self.file = open(file_name, method)
def __enter__(self): # the first thing that happens (entering the "with" block)
print("Enter")
return self.file
def __exit__(self, exc_type, exc_value, exc_tb): # the last thing that happens (handling the exception - closing the file)
print("Exit")
self.file.close()
if exc_type is Exception:
return True # telling Python that we handled the exception ("return True" suppresses it so no exception will arise)
with File("file.txt", "w") as f:
f.write("File")
raise Exception()
The @contextmanager
decorator from the contextlib
module provides an easier and cleaner way to create a context manager using a generator.
from contextlib import contextmanager
@contextmanager
def file(file_name, method):
file = open(file_name, method)
yield file
file.close()
with file("file.txt", "w") as f:
f.write("File")
The example below shows how context managers can be leveraged while working with our classes (the with
keyword also works with objects other than files).
class Counter:
def __init__(self, count=0):
self.count = count
class IncrementModifier:
def __init__(self, counter, amount):
self.counter = counter
self.amount = amount
def __enter__(self):
self.counter.count += self.amount # incrementing the count
def __exit__(self, exc_type, exc_val, exc_tb):
pass # no action needed on exit
counter = Counter()
print("Initial count:", counter.count)
with IncrementModifier(counter, 5): # using IncrementModifier to increase the count
print("Count after increment:", counter.count)
with IncrementModifier(counter, 3):
print("Count after second increment:", counter.count)
Context managers are also used to make sure a thread unlocks a thread lock.