Open In App

Reusable piece of python functionality for wrapping arbitrary blocks of code : Python Context Managers

Improve
Improve
Like Article
Like
Save
Share
Report

Context Managers are the tools for wrapping around arbitrary (free-form) blocks of code. One of the primary reasons to use a context manager is resource cleanliness. Context Manager ensures that the process performs steadily upon entering and on exit, it releases the resource. Even when the wrapped code raises an exception, the context manager guarantees the exit. So, without any procrastination, lets dive and acquire the new mantra for resource cleanliness without the code repetition.

Context Managers

Context Managers

Note: Knowledge about Decorators and Generators is appreciated

Context Manager Syntax

As we all know, the proper way to release the resource is by closing the resource after its use. And the most common practice followed to structure the close function is by using Exception Handling. Let’s look into the below code.




def main():
    try:
        file = open('sample.txt', 'r')
        data = file.read()
        print(data)
    finally:
        file.close()
  
main()


The finally clause ensures that the file will be closed irrespective of any situation. But this process results in code repetition. So, how you achieve the same using Context Manager without the boilerplate setup. Let’s go through the below sections one by one:

  • The with statement
  • The enter and exit method
  • Exception Handling

The with statement

with statement, the python built in function can be used as a context manager.




with open (‘filename’, ’r’)  as file:
    data = file.read()


This code performs the same functionality of what we developed using exception handling. Here the with statement is expected to return an object with two magic methods: __enter__ and __exit__. The result returned from __enter__ method is assigned to the variable mentioned after the `as` keyword.

It handles if an exception is raised within the wrapped code and most importantly the statement ensures that the resource is released. So, the programmers need not perform the close operation. Here is the full code.




def main():
    with open('sample.txt', 'r') as file:
        data = file.read()
        print(data)
  
main()


The enter and exit method

The __enter__ method returns an object, and then the returned value is assigned to the variable mentioned after the `as` keyword. Other than self argument it takes no other arguments.
On the other hand, the __exit__ method takes 3 positional arguments other than self argument. By default, these three arguments are none and are populated with information when the wrapped code raises an exception.

The below code explains the workflow of the Context Manager. Keep a note of the inside attribute and How the attribute value changes in different scenarios. Upon entering the context manager, the attribute value is set to true and on exit,  its value is set to false. Once the block of code within the with statement completes execution the __exit__ method gets invoked.




class ContextCheck(object):
      
    def __init__(self):
        self.inside = False
  
    def __enter__(self):
        self.inside = True
        return self
  
    def __exit__(self, exc_type, exc_instance, traceback):
        self.inside = False
  
  
cntCheck = ContextCheck()
print(cntCheck.inside)
  
with cntCheck:
    print(cntCheck.inside)
  
print(cntCheck.inside)


Output

False
True
False

Exception Handling

How the exception is handled using the Context Manager?
Using the __exit__ method, the context manager handles exceptions that are raised by the wrapped code. __exit__ method has 3 positional arguments:

  • Type of the Exception
  • An instance of the Exception
  • Traceback option

By default, all three values are None. When an __exit__ method receives an exception, the method can handle the exception in 2 different ways:

  • Re-raise the exception
  • Suppress the exception

Re-raise the exception

The __exit__ method can re-raise the exception by having a return statement that returns False. Let’s see the below code:




class BubbleExc(object):
    def __enter__(self):
        return self
  
    def __exit__(self, ex_type, ex_instance, traceback):
        if ex_instance:    
            print('Exception: % s.' % ex_instance)
        return False
  
with BubbleExc():
    1 / 0


Output

Exception: division by zero.
Traceback (most recent call last):
  File "C:\Users\Sonu George\Documents\Python\Context Managers\bubbleExc.py", line 11, in 
    1/0
ZeroDivisionError: division by zero

Suppress the exception

The __exit__ method can suppress the exception by returning true. See the below example for detailed understanding.




class SuppressExc(object):
    def __enter__(self):
        return self
  
    def __exit__(self, ex_type, ex_instance, traceback):
        if ex_instance:
            print('Suppressing exception: % s.'% ex_instance)
        return True
  
with SuppressExc():
    1 / 0


Output>

Suppressing exception: division by zero.

When you should write Context Managers

  • Releasing the Resources: Releasing the resources often ends up with boilerplate code. Using Context Managers, programmers can open the resource in the __enter__ method and close it in the __exit__ method and can reuse the functionality on demand.
  • Avoid Repetition: Exception Handling using except clause results in repetitive code. Using Context Manager, programmers can either raise or suppress the exception, and importantly it is defined in one place and hence can avoid repetition.

Writing Context Manager

Knowing is not enough; we must apply. Willing is not enough; we must do.
– Johann Wolfgang von Goethe

Implementing a Context Manager as a Class

A sample method that writes to a file using Context Manager. Here is the full code.




class FileClass(object):
      
    def __init__(self, f_name, func):
        self.f_obj = open(f_name, func)
          
    def __enter__(self):
        return self.f_obj
  
    def __exit__(self, type, value, traceback):
        self.f_obj.close()
  
with FileClass('sample.txt', 'w') as f_open:
    f_open.write('Congratulations, Good Work !')


Output

Congratulations, Good Work!

Closing a page using Context Manager

Here the page.close() will be called after completion of reading the page (upon exit from the with block).




from contextlib import closing
from urllib.request import urlopen
  
with closing(urlopen('https://www.geeksforgeeks.org/')) as page:
    for ln in page:
        print(ln)


Output


Using a Context Manager as a function Decorator and Generator

With the python contexlib module, you can implement context manager as decorator and the yield statement within the method gives the flexibility to use it as a generator. To understand this program, you should have prior knowledge about decorators and generators.




from contextlib import contextmanager
  
@contextmanager
def open_file(name):
    try:
        print('File Open')
        f = open(name, 'r')
        yield f
    finally:
        print('Close Operation')
        f.close()
  
  
  
def main():
    with open_file('sample.txt') as f:
        data = f.read()
        print(data)
main()


Here @contextmanager wraps the open_file method. When the function open_file gets invoked, the __enter__method of the context manager executes and pass the control to open_file method where the file opens and yield the file reference to callable with open_file('sample.txt') as f and the execution halts before executing the finally block.  

Once the code within the with statement is executed, it passes the control back to the open_file and the execution begins from where it is halted. In this case, the finally block executes and closes the file. Upon completion of execution, the __exit__ method is called and it handles any exception that is raised by the wrapped code.

Handling exception by checking instance of an Exception Class

In this case, we will check whether an exception is an instance of an exception class. Below we will create a subclass ‘TypeErrorSubclass’, which is derived from the ‘TypeError’ exception class, and using the raise statement we will raise the ‘TypeErrorSubclass’ exception. Here is the full code.




class TypeErrorSubClass(TypeError):
    pass
  
  
class ExceptionClass(object):
    def __enter__(self):
        return self
  
    def __exit__(self, ex_type, ex_instance, traceback):
        # Return True if there is no exception
        if not ex_type:
            return True
  
        # Return True if execution type is
        # equal or a subclass of TypeError
        if issubclass(ex_type, TypeError):
            print('Handling TypeError: % s' % ex_instance)
            return True
  
        return False
  
with ExceptionClass():
    raise TypeErrorSubClass('Type Error')


Output

Handling ValueError: Type Error

The __exit__ method checks if the exception raised is an instance of ‘TypeError’ class using issubclass(ex_type, TypeError), and if it’s an instance then suppress the exception by returning True.

Summary

Context managers ensure resources are released properly and the way it’s structured provides the flexibility of re-using the exception handling code at different places and hence avoid code repetition. Context managers are used for wrapping arbitrary blocks of code. However, a Context Manager can be used as a decorator, which avoids the need for writing a Context Manager and a Decorator separately for certain tasks.



Last Updated : 01 Aug, 2020
Like Article
Save Article
Previous
Next
Share your thoughts in the comments
Similar Reads