A Beginner’s Guide to Python Context Manager

 When we watch Youtube videos or any tutorial about data analysis, we see that the context manager is mainly used in reading or writing some texts from a file so that we don’t have to write the code for closing the file whenever our task is finished. But this is not the only use-case of this python context manager. We can use context managers in more relaxed ways. For example, we can easily interchange the working directory using a context manager. We can also access a database efficiently by this. In this article, we will discuss everything about it and see some examples that will help us use them properly in your project. So what are we waiting for? Let’s get started.

This article was published as a part of the Data Science Blogathon.

What is a Context Manager?

 A context manager is actually a resource manager in python. For example, when we read or write some content from a file, we don’t care about closing the file. If we are using one file, then this is not a big issue. But when it comes to multiple files, this will cause problems as we are not releasing the resources we are no longer using. This is also true for database access. After a database is used, we should close the database. Otherwise, it can slow down our code. So, the context manager helps us maintain the resource adequately. Below is an example of a context manager used for accessing the file.

with open(“file.txt”,”r”) as f:
        print(f.read())

Here we are trying to open the file in read mode using python context manager and reading the first line from the file. The file will automatically close when we leave the context.

_{Creating Context Managers}

In the previous example, we used the open() function as a python context manager. There are two ways to define a context manager – one is class-based and the other is function-based. Let’s see how we can create it using class.

Python Code:

class FileManager():
    def __init__(self, filename, mode):
        self.filename = filename
        self.mode = mode
        self.file = None
    def __enter__(self):
        self.file = open(self.filename, self.mode)
        return self.file
    def __exit__(self, exc_type, exc_value, exc_traceback):
        self.file.close()

For the class-based, we have to add two methods: enter () and exit(). In the enter() method, we define the things we want to do in the context. Here we open a file. In the exit() method, we have to write things that will happen after leaving the context. In this case, we will close the file after leaving the context. Now let’s see what a function-based python context manager looks like.

import contextlib
@contextlib.contextmanager
def my_context():
      # add any set up code you need
      yield
      # add any teardown code you need

For a function-based context manager, you define a function with the @contextlib.contextmanager decorator. Within that function, you place the code you want to execute within the context. Following that, you use the yield keyword to indicate that this function serves as a python context manager. After yield, you can include your teardown code, specifying actions to take when exiting the context. This teardown code ensures proper resource cleanup and is executed when leaving the context.

_Examples

We discussed the theoretical part of the context manager. Now it is time to see some examples. Let’s see how we can access the database.

import psycopg2

@contextlib.contextmanager
def database(url):
  # set up database connection
  db = psycopg2.connect(url)
  yield db
  # tear down database connection
  db.disconnect()

You have already understood what is happening in the function. This context manager will give access to a database in its context. When we leave the context, the connection to the database is ended. You don’t have to return value whenever creating a context manager.

Now here is another example. While using Google Colab, if the data size is too large, we attach our Google Drive to the notebook and create a folder for the data. Then we download the data in that folder and want to return it to our working directory. This we can do easily with a context manager.

@contextlib.contextmanager 
def in _ dir(path): 
# save current working directory 
old _ dir = os.getcwd() 
# switch to new working directory
 os.chdir(path) 
yield 
# change back to previous 
# working directory 
os.chdir(old _ dir)

See the above example. We don’t return any value from the “yield” keyword. I hope you now understand how to use context manager.

Managing Multiple Files

Sometimes it happens that we need to open multiple files. So what should we do? Of course, we have to use multiple context managers. Here is an example of copying the content of a file.

with open(“file.txt”, “r”) as f:
	contents = f.read()
with open(“copy.txt”, “w”) as cf:
	cf.write(contents)

This will work, but if we have a large file to copy, it will cause pain for us. That will make the code slower. You can also face device freezing as that process takes a lot of resources. So what should we do? If we can copy the file line by line, that doesn’t cause any problem. To do this, we have to open two files simultaneously. See the below example.

with open(“file.txt”, “r”) as f:
       with open(“copy.txt”, “w”) as cf:
	   for line in f:
		cf.write(line)

This is called the nested context. This method is fine when you have to open two files or three files. But when the number of files is much more, the code becomes messier and you have to face difficulties while editing the code. In Python 3.10, they introduced a new method “parenthesized context” in which you can write multiple contexts in one line using brackets. No need to use looping.

with (open(“file.txt”, “r”) as f, open(“copy.txt”, “w”) as cf):
        for line in f:
	  cf.write(line)

The above code is much cleaner. For me, this is the best option for working on multiple files. Now it may happen that you will prefer others. That’s totally up to you.

Conclusion

Using context manager in your code ensures the proper utilization of the resource. Resource utilization is not a big deal when you are working on a small project, but when it comes to large projects, like reading a very large file that we discussed previously, context manager comes in handy. In this article, we learned-

• The basics – what and how.
• The types of context managers, according to their implementation – are class-based and function-based.
• How to implement the different types of context managers.
• How to handle multiple files using nested context manager and parenthesized context manager.

Thank you for reading this article till the last. I hope you enjoyed this article. If there is any issue, please leave a comment. And don’t forget to check out my other articles in AnalyticsVidhya.

Frequently Asked Questions

The media shown in this article is not owned by Analytics Vidhya and is used at the Author’s discretion. 

Subhradeep

Recently pursuing M.Tech in Artificial Intelligence and love to do anything about Data Science, Machine Learning, and AI. I also like to share my knowledge through Blogs. Ask me anything about Data Science, Machine Learning, and AI at [email protected].