Guide To Python Context Managers: Beyond Files

Context managers are one of those inherently Pythonic features that aren't available in most languages. They provide a convenient syntax when your code has to open and close connections or manage limited resources.
python context managers

Python context managers are a nifty resource management tool that provides a simple syntax for a powerful construct. When working with threads, database connections, files, or any limited resources, it’s essential to release them after use. Not releasing these resources leads to resource leakage and can cause the program to act erratically or even crash. Generally, the allocating and releasing of resources is achieved by using try/finally, utilizing the resource in the try block and releasing it in the finally block. Python, however, has a better, more convenient way. And even if this is the first time you’ve come across the term “context managers”, you’re probably familiar ‘with’ it.  This article will explore what context managers are, how they work, how to implement your own and then some use-cases. 

What is a Context Manager?

Python provides the with statement that enables the developer to implement the concept of a runtime context. This can be broken down and understood using the analogy of a local variable inside a function. The local variable is created when the function is called and is, ideally, automatically deleted when the function has finished execution. In this case, we are releasing memory when we’re done with it. Similarly, in the case of context managers, the resource is locked or marked as ‘in use’ when the interpreter enters the context manager block and released as soon as it’s exited.

Let’s see the trusty with statement in action and compare it with its more verbose try/finally alternative:

Subscribe to our Newsletter

Join our editors every weekday evening as they steer you through the most significant news of the day, introduce you to fresh perspectives, and provide unexpected moments of joy
Your newsletter subscriptions are subject to AIM Privacy Policy and Terms and Conditions.
 with open('the-tales-of-beadle-the-bard.txt', 'w') as f:
     f.write("If Harry Potter speaks Parseltongue, can he code in Python?") 

is equivalent to:

     f = open('the-tales-of-beadle-the-bard.txt', 'w')
     f.write("If Harry Potter speaks Parseltongue, can he code in Python?")

How to Write Custom Python Context Managers

There are two ways of implementing custom context managers in Python:

As a Class

Context managers can be implemented as classes using the two dunder methods __enter__ and __exit__.  The with statement calls the __enter__ method, and the value returned by this method is assigned to the identifier in as clause of the statement. After the __enter__ method is called, the code inside the with block is executed, and then the __exit__ method is called with three arguments – exception type, exception value and traceback. All three arguments are None if the code inside the with block doesn’t raise any exceptions. If there’s is an exception it can be handled in the __exit__ method using the three arguments. Let’s illustrate this by creating a timed file handling context manager.

 from time import perf_counter, sleep

 class TimedFileHandler():
     def __init__(self, file_name, method):
         self._file_name = file_name
         self._method = method
         self._start = 0.0
     def __enter__(self):
         print(f"Entered file: {self._file_name}, in mode: {self._method}")
         self._start = perf_counter()
         self._file = open(self._file_name, self._method)
         return self._file
     def __exit__(self, exc_type, value, traceback):
         if exception_type is not None:
             print("Wrong filename or mode")
         end = perf_counter()
         print("Exited file")
         print(f"Time taken: {end - self._start}")

 with TimedFileHandler('the-tales-of-beadle-the-bard.txt', 'w') as f:
     f.write('C is simple.\n *(*(*var)())[7]X;')
 print(f"File closed?: {f.closed}") 
 Entered file: the-tales-of-beadle-the-bard.txt in mode: w
 Exited file
 Time taken: 5.001780649000011
 File closed?: True 

The same outcome can be achieved using by creating an object of the class:

 handler = TimedFileHandler('the-tales-of-beadle-the-bard.txt', 'w') 
 f = handler.__enter__()  # allocating
 f.write('C is simple.\n *(*(*var)())[7]X;') # executing body 
 handler.__exit__(None, None, None)  # releasing 
As a Generator

The contextmanager decorator from the contextlib module provides a more convenient way of implementing context managers. When the contextmanager decorator is used on a generator function, it automatically implements the required __enter__ and __exit__ methods and returns a context manager instead of the unusual iterator.

 import contextlib
 def timed_file_hanlder(file_name, method):
     start = perf_counter()
     file = open(file_name, method)
     print(f"Entered file: {self._file_name}, in mode: {self._method}")
     yield file
     end = perf_counter()
     print("Exited file")
     print(f"Time taken: {end - start}")

 with timed_file_hanlder("the-tales-of-beadle-the-bard.txt","w") as f:
        f.write("*Code written last month*\ I have no memory of this place.")
 print(f"File closed?: {f.closed}")
 Entered file: the-tales-of-beadle-the-bard.txt in mode: w
 Exited file
 Time taken: 5.005866797999943
 File closed?: True 

Context Managers In Action

Django provides many useful context managers, such as transaction.atomic that enables developers to guarantee the atomicity of a database within a block of code. If the code in the with block raises an exception, all the changes made by the transaction are reverted.

 from django.db import transaction
 with transaction.atomic():
     # This code executes inside a transaction.

The Session class of the request module implements the __enter__ and __exit__ methods and can be used as a context manager when you need to preserve cookies between requests, want to make multiple requests to the same host or just want to keep TCP connection alive.

 import requests
 with requests.Session() as sess:
     sess.request(method=method, url=url) 

If you’ve read our previous articles on multithreading and multiprocessing, you would remember that we used the ThreadPoolExecutor and ProcessPoolExecutor context manager to automatically create and run threads/processes.

  with concurrent.futures.ProcessPoolExecutor() as executor:
      secs = [5, 4, 3, 2, 1]
      pool = [executor.submit(useless_function, i) for i in secs]
      for i in concurrent.futures.as_completed(pool):
          print(f'Return Value: {i.result()}')
  end = time.perf_counter()
  print(f'Finished in {round(end-start, 2)} second(s)')  

You can find the above code in a Colab notebook here.

Last Epoch

Python’s with statement provides a more convenient syntax when your code has to open and close connections or manage limited resources. There are more benefits to using the with statement than just a more concise syntax. The first advantage is that the whole allocation-release process happens under the control of the context manager object, making error prevention easier. Another reason to use it is that the with block makes it easier for developers to discern where the connection/resource is in use or can be used. Furthermore, it acts as a refactoring tool that moves the code for setup and teardown of any pair of operations to one place – the method definitions. 

Context managers are one of those inherently Pythonic features that aren’t available in most languages. They provide a simple and elegant solution to a plethora of problems, even beyond resource management. I hope this article helped you better understand context managers so you can start using them to write better code. 

If you want to learn more about context manager and the contextlib module, refer to the following resources:

Aditya Singh
A machine learning enthusiast with a knack for finding patterns. In my free time, I like to delve into the world of non-fiction books and video essays.

Download our Mobile App


AI Hackathons, Coding & Learning

Host Hackathons & Recruit Great Data Talent!

AIM Research

Pioneering advanced AI market research

Request Customised Insights & Surveys for the AI Industry


Strengthen Critical AI Skills with Trusted Corporate AI Training

Our customized corporate training program on Generative AI provides a unique opportunity to empower, retain, and advance your talent.

AIM Leaders Council

World’s Biggest Community Exclusively For Senior Executives In Data Science And Analytics.

3 Ways to Join our Community

Telegram group

Discover special offers, top stories, upcoming events, and more.

Discord Server

Stay Connected with a larger ecosystem of data science and ML Professionals

Subscribe to our Daily newsletter

Get our daily awesome stories & videos in your inbox