How to Create a Custom Context Manager in Python OOP: Python OOP Complete Course (Part 19)

Learn what context managers special methods in Python OOP are and how to override them.

Before we start let me tell you that:

This article is a part of The Complete Course in Object-Oriented Programming in Python which you can find it here.
All resources are available in the “Resources” section below.

Introduction

When you hear the term context managers, the keyword that you have to remember is with

As a reminder the with keyword migh be used to open a file for reading or writing as follows:

with open('file_path', 'w') as file: 
    file.write('hello world !')

Usually, we use with keyword with files to guarantee closing the file after completing all the required processing.

So, how to make your object able to use with keyword in Python? Let us see together.

Table of Contents

What are Context Manager’s Special Methods?
Available Context Manager’s Special Methods
How to Override Context Manager’s Special Methods?

1. What are Context Manager’s Special Methods?

In general, context managers allow setup and cleanup resources for objects when their creation is wrapped with a with statement (For example, databases are one of the resources that need to be cleaned up after you finish using them. This cleaning happens by closing the connection).

In particular, context managers’ special methods are the methods that are used to define the behavior of the context manager. In other words, it defines what to do when you enter the block of code of that statement and what to do when you leave this code block.

Now, let us check what are the available context manager special methods.

2. Available Context Manager’s Special Methods

As has been mentioned before, you have to define the behavior at the beginning of the code block of with statement and the end of this code block.

Therefore, there are two context managers methods:

__enter__(self): it defines what the context manager should do at the beginning of the block created by with statement.
__exit__(self, exception_type, exception_value, traceback): it defines what the context manager should do after its block has been executed or terminated.

Now let us see a concrete example to learn how to override them.

3. How to Override Context Manager’s Special Methods?

First, let us take an example that shows why context managers are useful.

Assume you have a folder, and its name dataset, and inside this folder you have the file dummy.txt.

Let us open this file in two cases.

Without using (context manager).

Output:

Hi I am Dummy File
True

The previous example has been executed without any errors, and as you have seen the file.closed returned True which means the file is closed.

Now, what will happen if an error happens before you close the file?

Output:

IndexError                                Traceback (most recent call last)
<ipython-input-2-9755d8bd0574> in <module>
      4 content = file.read()
      5 a = [1, 2]
----> 6 b = a[10]
      7 file.close()
      8 print(content)

IndexError: list index out of range

In the previous example, you tried to access the 10th element from the list that has only two elements, so you got an error.

Check if the file is closed.

print(file.closed)

Output:

False

You got False because an error occurred before calling the function file.close() and the file hasn’t been closed.

Now let us try the previous two examples with the context managers.

Output:

Hi I am Dummy File
True

As you can see, it has run and the file has been closed because the context manager closed the file beyond the scene.

Let us check what will happen if there is an error has happened before you close the file.

Output:

IndexError                                Traceback (most recent call last)
<ipython-input-6-7d029e3dc81d> in <module>
      3 with open(file_path) as fil:
      4     a = [1, 2]
----> 5     b = a[10]
      6     content = fil.read()
      7

IndexError: list index out of range

As you can see, you got an error.

Check if the file is closed.

print(fil.closed)

Output:

True

You got True even though an error happened.

So, the context manager helps you to clean up the resources and close the files even if there is an error.

Now let us define the Fileclass which simulates open() function.

Assume the class Filehas three attributes (file, file_path, and mode), and you want to make your class as a context manager, so you should override the __enter__() method and the __exit__() method.

Output:

Hi I am Dummy File

The file has been closed

As you can see, the class opened the file, and then it closed it without any problems.

Check if the file is closed.

print(f.closed)

Output:

True

You got True which means that your context manager class File closed the file after it exited from the body of the with statement.

Check what will happen when an exception occurs within with scope.

with File(file_path, 'r') as f:
    cont = f.read()
    a = [1,2]
    b = a[10]
    print(cont)
    print()

Output:

The file has been closed

---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-11-61758e09ee00> in <module>
      2     cont = f.read()
      3     a = [1,2]
----> 4     b = a[10]
      5     print(cont)
      6     print()

IndexError: list index out of range

In the previous example, you tried to access the 10th element from the list that has only two elements, so you got an error.

Check if the file is closed.

print(f.closed)

Output:

True

You got True which means the file is closed even though an error has happened.

Now, assume you want to access the type of the error so you can handle it. To do that you should update __exit__() method.

Output:

The file has been closed
exception_type: <class 'IndexError'>
exception_value: list index out of range
traceback: <traceback object at 0x0000000004DD3D88>

---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-14-565acbd5e564> in <module>
     22     cont = f.read()
     23     a = [1,2]
---> 24     b = a[10]
     25     print(cont)
     26     print()

IndexError: list index out of range

As you can see, when an error occurs inside the context manager scope, you get all the details about this error so you can handle it easily according to the use case that you are working on.

Now, let us summarize what we have learned in this article.

Context Manager’s Special Methods: are the methods that are used to define the behavior of the context manager at the beginning and at the end of the code block of with statement.
enter() method: defines what the context manager should do at the beginning of the block created by thewith statement.
exit() method: defines what the context manager should do after the block has been executed or terminated.

P.S.: A million thanks for your time reading my story. Before you leave let me mention quickly two points:

First, to get my posts in your inbox directly, would you please subscribe here, and you can follow me here.
Second, writers made thousands of $$ on Medium. To get unlimited access to Medium stories and start earning, sign up now for Medium membership which only costs $5 per month. By signing up with this link, you can directly support me at no extra cost to you.

The Complete Course in Object-Oriented Programming in Python

undefined

To get back to the previous article, you can use the following link:

Part 18: How to Create a Customized Sequence Class

To move on to the next article, you can use the following link:

Part 20: How to Create Callable Objects in Python

Resources:

GitHub here.