This post is inspired from a talk at python pune january meetup by Pradhavan.
What are Context Managers?
Here's what Python's official documentation says:
A context manager is an object that defines the runtime context to be established when executing a with statement. The context manager handles the entry into, and the exit from, the desired runtime context for the execution of the block of code. Context managers are normally invoked using the with statement, but can also be used by directly invoking their methods. – Python Docs
But that's just covering all the bases. Let's understand a simplified version.
Programmers work with external resources from time to time, like files, database connections, locks etc. Context managers allow us to manage those resources by specifying:
- What to do when we acquire the resource, and
- What to do when the resource gets released
Why do we need Context Managers?
Consider the following example:
for _ in range(100000):
file = open("foo.txt", "w")
files.append(f)
file.close()
Notice that we're calling the close() method to ensure that the file descriptor is released every time. If we didn't do that, our OS would run out of its allowed limit to open file descriptors eventually.
However, we write a more pythonic version of the above code using context manager:
for _ in range(100000):
with open("foo.txt", "r") as f:
files.append(f)
Here open("foo.txt", "r") is the context manager that gets activated using the with statement. Notice that we didn't need to explicitly close the file, the context manager took care of it for us. Similarly there are other predefined context managers in Python that makes our work easier.
Can we define our own Context Manager?
Yes. There are two ways to define a custom context manager:
- Class based definition.
- Function based definition.
Class Based Context Managers
Let's continue with our file example and try to define our own context manager which will emulate open().
files = []
class Open():
def __init__(self, filename, mode):
self.filename = filename
self.mode = mode
def __enter__(self):
self.open_file = open(self.filename, self.mode)
print("Enter called")
return self.open_file
def __exit__(self, *args):
print("Exit Called")
self.open_file.close()
for _ in range(100000):
with Open('foo.txt', 'w') as f:
files.append(f)
- The
__enter__method tells us what to do when we acquire the resource, i.e., giving us a file object. - The
__exit__method defines what to do when we're exiting the context manager, i.e., closing the file. - You can see how both
__enter__and__exit__are called with every loop.
Handling Errors
How do we handle FileNotFoundError with python's open()
try:
with open("foo.txt", "r") as f:
content = f.readlines()
except FileNotFoundError as e:
print("Hey, file isn't there. Let's log it.")
Such a basic error handling code that needs to be every time you open a file. Let's try to DRY it with our custom context manager.
class Open():
def __init__(self, filename, mode):
self.filename = filename
self.mode = mode
def __enter__(self):
print("Enter called")
try:
self.open_file = open(self.filename, self.mode)
return self.open_file
except FileNotFoundError as e:
print("Hey, file isn't there. Let's log it.")
def __exit__(self, exc_type, exc_value, exc_traceback): #notice the parameters
print("Exit Called")
if(exc_type is None):
self.open_file.close()
return True
else:
return True
with Open("foo.txt", "r") as f:
content = f.readlines()
Changes in __exit__
-
exc_typeis type of error Class which you'll get while handling errors in__enter__(AttributeError in this case). -
exc_valueis the value of the error which you'll get while handling errors in__enter__. -
exc_tracebackis the traceback of the error which you'll get while handling errors in__enter__. - We're returning
Trueto suppress the error traceback (not to be confused withexc_tracebackparameter).
Another Real World Example
class DatabaseHandler():
def __init__(self):
self.host = '127.0.0.1'
self.user = 'dev'
self.password = 'dev@123'
self.db = 'foobar'
self.port = '5432'
self.connection = None
self.cursor = None
def __enter__(self):
self.connection = psycopg2.connect(
user=self.user,
password=self.password,
host=self.host,
port=self.port,
database=self.db
)
self.cursor = self.connection.cursor()
return self.cursor
def __exit__(self, *args):
self.cursor.close()
self.connection.close()
Function Based Context Managers
Function based context management is done by using a lib called contextlib, through which we can change a simple generator function into a context manager. Here's what a typical blueprint looks like:
from contextlib import contextmanager
@contextmanager
def foobar():
print("What you would typically put in __enter__")
yield {}
print("What you would typically put in __exit__")
with foobar() as f:
print(f)
-
contextmanagerdecorator is used to turn any generator function into a context manager. -
yieldwork as a separater between__enter__and__exit__parts of the context manager.
Handling files
from contextlib import contextmanager
@contextmanager
def open_(filename, mode):
print("SETUP")
open_file = open(filename, mode)
try:
print("EXECUTION")
yield open_file
except:
print("Hey, file isn't there. Let's log it.")
finally:
print("CLEAN-UP")
open_file.close()
with open_("somethign.txt", "w") as f: #notice the mode
content = f.readlines() #you cannot read on write mode
We wrap yield in a try block because we don't know what the user is going to do with the file object. They might try to use it in a way that it's not intended to (as shown above).
Database Connections
from contextlib import contextmanager
@contextmanager
def database_handler():
try:
host = '127.0.0.1'
user = 'dev'
password = 'dev@123'
db = 'foobar'
port = '5432'
connection = psycopg2.connect(
user=user,
password=password,
host=host,
port=port,
database=db
)
cursor = connection.cursor()
yield cursor
except:
print("Hey, file isn't there. Let's log it.")
finally:
cursor.close()
connection.close()
Resources
We have just covered just an introduction to context managers, but I feel that it's just the tip of the iceberg and there are many interesting use cases for it. Here are some interesting links that I found:
Top comments (0)