Python

Python Context Managers - A Better Way to Manage Resources

Context managers allow us to properly manage system resources by specifying what we want to set up and tear down while working with objects. A system resource can be a file handler, a database connection, a network socket, a lock, etc. They look strange when someone sees them the first time, but one can add them into their programming arsenal as a powerful weapon. 

If you are a new Python programmer then you might have already written code to process file content. The most commonly used code for file processing would look like this:

fp = open('hello.txt', 'w')
fp.write('this is my content')
fp.close()

The above code opens a text file, writes some content in it, and then closes it. Every file which is opened will need to be closed to avoid any memory leakage issues that can also cause your program to crash. The above code looks a bit incomplete as it lacks logic to handle errors that may occur between the file open and close statements. Here is a more refined code to achieve the same task:

fp = open('hello.txt', 'w')
try:
   fp. write('this is my content')
finally:
   fp.close()

Wrapping the code by a try/finally block has made the code less error-prone and ensures that the opened file will always be closed automatically. But this has come at the cost of some boilerplate code. It’s the right time to use a context manager using a with statement. A with statement is the most seen form of a context manager in Python and below is a pseudo-code of that:

with EXPR :
   CODE-BLOCK

Many context managers return an object when they are called and you can assign that object to a new target variable name using an as clause.

with EXPR as VAR:
  CODE-BLOCK

Below is an equivalent code of the above file processing task using a context-manager:

with open('hello.txt', 'w') as fp:
   fp.write('this is my content')

In the code above, a with statement ensures that file resources will be closed once all statements defined within that block are complete or even if the program failed because of an error in the block. We no longer need to worry about closing the file, as the context manager will do that itself. It has also eliminated much of the unnecessary boilerplate code.

 

More code examples of Context Managers :

There are many places where you can see context managers are used to manage resources. Database connections, thread-locking, and network socket are some of the common areas. Let us see some more code examples:

Thread locking example:

# Code without a context manager
import threading
lock = threading.Lock()
lock.acquire()
try:
   # some operations here
finally:
   lock.release()

A more preferred code would be:

# code with a context manager
import threading
lock = threading.Lock()
with lock:
   # some operations here

Network socket code:

# code sample to context manager to work with socket
import socket
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as transfer:
   # socket connection and other operations here

 

Custom Context Managers:

So far we have seen the in-built context managers which come with default libraries but we don’t know how they work and manage resources behind the scenes. Python allows us to write our own context-managers to handle resources, and this gives us a great opportunity to understand the techniques of how they work.

There are a couple of different ways to write your own custom context-managers by using a class or by using a function with a decorator.

 

1. Writing a custom context manager using a class:

Let us first see a class-based context-manager as it gives more clarity to understand them. The basic prerequisite to turn a Python class into context-manager requires defining two magic methods __enter__ and __exit__.

class CustomContextManager():

    def __init__(self, args):
        pass

    def __enter__(self):
        pass

    def __exit__(self, exc_type, exc_val, traceback):
        pass

__enter__ and __exit__ are set-up and tear-down methods for writing custom context managers. __init__ method is to pass arguments to class or to set some class attributes which can be used by the other two methods. Let us take the same above file processing task example to write our custom context manager by defining a FileManager class.

class FileManager():
   def __init__(self, file_name, mode):
     self.file_name = file_name
     self.mode = mode

   def __enter__(self):
     self.fp = open(self.file_name, self.mode) 
     return self.fp

   def __exit__(self, exc_type, exc_val, traceback):
     self.fp.close()

Let’s use this context manager using the same with statement and then walk through it to point out exactly what is going on there.

with FileManager('hello.txt', 'w') as f:
   f.write('this is my content')

  • __init__ method takes 2 arguments file_name and mode (can be read/write/append) and sets them as class properties.
  • FileManager class is being instantiated using a with statement (a context manager property) so it calls the __enter__ method automatically (after the __init__ method). You can see the __enter__ method opens the file and returns the file (descriptor) object and the object it assigns to an ‘f’ variable using as clause to perform actions on the file.
  • __exit__ method gets invoked automatically just before the program jumps out of the with block. This method has logic to close the file so we need not worry about closing it anymore.

I hope now you are clear about how to write your own context manager using a class and also got a good understanding of the techniques used by context managers to manage resources.

 

2. Writing a custom context manager using a function (with contextlib):

Python has a module named contextlib which allows writing context managers using a function. This requires you to decorate your function with an @contextmanager decorator which is a utility of contextlib. I am going to take the same file processing example from above to show you a function-based context manager.

from contextlib import contextmanager

def open_file(file_name, mode):
    fp = open(file_name, mode)
    try:
        yield fp 
    finally:
        fp.close()

This code does everything equivalent to our above class-based context manager, but it is a little less than that. Though this may look difficult to understand if you are not familiar with decorators and generators in Python. I don’t want to spend time teaching you about them, as this article is not about them.

In Python, a yield statement suspends function execution and returns a value to its caller, and retains state to enable function resuming where it was left off. A function that uses a yield statement to generate value rather than a return is known as a generator function.

If you look at our function, then you can see there is one yield statement. The code before the yield statement is equivalent to what our __enter__ method of the class was (a setup of context manager). Everything after the yield statement is equivalent to what was in the __exit__ method (a teardown of the context manager). Let's see the code sample of our generator function with a with statement.

with open_file('hello.txt', 'w') as f:
   f.write('this is my content')

That's it!

Share post

  •  
  •  
  •  
  •  

Tilak S.

Technology freak, Open Source lover. Someone trying to understand many things. Wants to make a difference. Life liver and Peripatetic.

Other posts you might like