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.