Contents

Context Manager Protocol

  • Allows for some common resource management by making the code more expressive and avoid resource leaks.

  • Ensures that a used resource (file, database etc.) is cleaned up whenever a code that uses it is done, even if exceptions are thrown.

  • Provides a syntactic sugar for try/finally blocks.

Consider below scenario, where you need to access a file, do some operation and close them.

file = open('hello.txt', 'r')
try:
    print(file.readline())
finally:
    file.close()

Context Manager Example

The above code can be re-written using with statement:

with open('hello.txt', 'r') as file:
    print(file.readline())

Context Manager Example

As we can see, this implementation is less verbose and does the same thing as the code above it:

  • Open the file

  • Do some operation

  • Close it (In case of any error, it tries to close it.)

Implementing your own Context Manager using a Class

  • At bare minimum, we need to implement two [dunder] methods: __enter__ and __exit__.

  • Lets implement our own file Context Manger as above:

class FileManager:
    def __init__(self, file, mode):
        self._file = open(file, mode)
    
    def __enter__(self):
        return self._file
    
    def __exit__(self, type, value, traceback):
        self._file.close()

Now we can use with statement as below:

with FileManager(file="hello.txt", mode="r") as file:
    print(file.readline())

Context Manager Example

  • The interpreter calls the __enter__ method when execution enters the context of with statement.

  • Similarly __exit__ method is called when execution leaves the context again.

  • The __exit__ method accepts three arguments:

    • type

    • value

    • traceback

  • They are the required parameters as part of the Context Manager protocol.

  • In case of exception, The interpreter (Python) will pass the type, value, and the traceback of the exception to the __exit__ method.

Note: The interpreter will store __exit__ method in advance, so that it can call it while leaving the context.

with FileManager(file="hello.txt", mode="r") as file:
    file.exception_should_occur()

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-5-b261b03a5a0c> in <module>
      1 with FileManager(file="hello.txt", mode="r") as file:
----> 2     file.exception_should_occur()

AttributeError: '_io.TextIOWrapper' object has no attribute 'exception_should_occur'

Let’s handle the exception in __exit__ method:

class FileManager:
    def __init__(self, file, mode):
        self._file = open(file, mode)
    
    def __enter__(self):
        return self._file
    
    def __exit__(self, type, value, traceback):
        print(
            "Exception handled:"
            f"\n\ttype: {type}, "
            f"\n\tvalue: {value}, "
            f"\n\ttraceback: {traceback}"
        )
        self._file.close()
        return True

with FileManager(file="hello.txt", mode="r") as file:
    file.exception_should_occur()

Exception handled:
	type: <class 'AttributeError'>, 
	value: '_io.TextIOWrapper' object has no attribute 'exception_should_occur', 
	traceback: <traceback object at 0x7f6df5ec7d00>

Implementing your own Context Manager using a generators

  • Python provides contextlib module to implement context managers using decorators and generators

  • Consider, below example:

from contextlib import contextmanager

@contextmanager
def file_manager(name, mode="r"):
    file = open(name, mode)
    try: 
        yield file
    finally:
        file.close()

contextmanager in action:

with file_manager('hello.txt') as file:
    print(file.readline())

Context Manager Example

  • The above implmentation looks very simple as compared to our class bases solution.

  • We are using contextmanager decorator to define a generator function that automatically supports with statement

  • In our example:

    • file_manager() function first acquires the file resource

    • It then temporarily suspends the execution and yields the file resource so it can be used.

    • When the execution leaves the with context, the generator continues to execute the remaining part of the code which usually contains the clean-up steps.

  • This approach requires some knowledge of decorator and generator concepts as opposed to class based implementation.

Some more examples:

# using class
import sqlite3

class SqliteDatabase:
    def __init__(self, database):
        self._database = database
        self._connection = sqlite3.connect(self._database)

    def __enter__(self):
        return self._connection
    
    def __exit__(self, type, value, traceback):
        self._connection.close()

def get_contacts():
    with SqliteDatabase('contacts.db') as connection:
        cursor = connection.cursor()
        for row in cursor.execute("SELECT * from contact"):
            yield row


for contact in get_contacts():
    print(contact)

(1, 'Sajal Shrestha', 'Male', 'Kathmandu, Nepal')
(2, 'Prashant Paudel', 'Male', 'Banepa, Nepal')
(3, 'Atul Shrestha', 'Male', 'Okhaldhunga, Nepal')
(4, 'Resa Manandhar', 'Female', 'Kathmandu, Nepal')

# using contextmanager
from contextlib import contextmanager
import sqlite3

@contextmanager
def sqlite_manager(database):
    connection = sqlite3.connect(database)
    try:
        yield connection
    finally:
        connection.close()

def get_contacts():
    with SqliteDatabase('contacts.db') as connection:
        cursor = connection.cursor()
        for row in cursor.execute("SELECT * from contact"):
            yield row

for contact in get_contacts():
    print(contact)

(1, 'Sajal Shrestha', 'Male', 'Kathmandu, Nepal')
(2, 'Prashant Paudel', 'Male', 'Banepa, Nepal')
(3, 'Atul Shrestha', 'Male', 'Okhaldhunga, Nepal')
(4, 'Resa Manandhar', 'Female', 'Kathmandu, Nepal')
Iterators Regular Expressions