Most of the Python programmers must have, either knowingly or unknowingly have dealt with a Context Manager.
Python provides such ease of using abstractions that most of the time we overlook the clever feature thinking it’s a bit trivial abstraction. Context Managers are no exception in that case.
If we look into the Python Docs this is what we would find.
A context manager is an object that defines the runtime context to be established when executing a with statement. - Python Docs
This might be a bit much for most of us, so in simple terms:
Objects that help with control flow of the with statement - Luciano Ramalho, Fluent Python
The most common Context Manager that every Python programmer uses very frequently for managing files, is with the open function
with open('StarWars.txt') as f:
f.write('This is the way!')Here a temp context is set when the
openfunction returns an IO object herefYou write to the file with the IO object
fWhen all things are done, that temp context is reliably closed. Somewhat like a setup/teardown process.
You must think that it is fine and all but why can’t I use a simple try/finally block in a similar manner. Well technically can do that too.
try:
file = open('StarWars.txt')
file.write('This is the way!')
except FileNotFoundError as e:
logging.error('File not found', exec_info=True)
finally:
file.close()But then you would be defeating the whole purpose of a Context Manager.
Context Managers exists to control a with statement and the with statement was designed to simplify the try/finally pattern.
Besides that, a Context Manager can help in improving code readability of the code and abstracting complex logic with a single withas control flow.
The withas control flow
Now that we have some basic idea of what a Context Manager is? and why do they exist? Let's come back to our file example and take a bit of a deeper dive to see how the withas control flow is executed.
with open('StarWars.txt') as f:
f.write('This is the way!')Withstatement is called with an expression.The result of that expression is the value of the variable after
asin thewithasstatement herefSome work is done inside the
withblock withfinside thewithblock.Finally when the control flow exits the
withblock the file is closed.
In other cases this closing of resource can also be releasing a resource that was being used or a previous state is being restored that was being changed inside the with block.
Writing Your Own Context Manager
Context Manager can be written in two ways either create your own class or use the Contextlib module from the standard lib.
Creating a class with magic methods:
class FileManager:
def __init__(self, filename, mode='r'):
self.filename = filename
self.mode = mode
def __enter__(self):
self.open_file = Open(self.filename, self.mode)
return self.open_file
def __exit__(self, exception, exception_type, exception_traceback):
self.open_file.close()
>> with FileManager('StarWars.txt', 'w') as f:
f.write('This is the way!')FileManager class tries to mimic the open function and exposes the two most important dunder methods __enter__ and __exit__ which are responsible for creating a temp context, doing the work, and restoring the state or object.
When the
withblock is called, dunder__enter__is called with just the self as an argument.In the enter method, an IO object i.e
self.open_fileis returned. This returned value is off, that is the variable that is used after theasstatement.fis used to write to the file and when thewithcontrol flow is moved out of the block and the__exit__
On a successful exit, all the three arguments of the dunder exit will have a NONE value. For some reason, if you want to suppress it, you should return True as anything returned besides True raises exception.
Now let’s look at somewhat real-world examples of Class-based Context Manager.
~ Real World Example
class DatabaseHandler:
def __init__(self):
self.host = 'localhost'
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, exc_type, exc_value, traceback):
self.cursor.close()
self.connection.close()
DatabaseHandler is the Postgres handler that we can import to other parts of the program.
__enter__method creates a DB connection and returns a cursor object__exit__method checks for exceptions, raises if any, and closes the connection reliably.
When using DatabaseHandler we don't have to worry much about opening/closing the connection or even managing exceptions. Which in return makes our code more readable.
class Editor:
def get_articles(self, start_date, end_date):
sql = 'Select query to get total articles'
with DatabaseHandler() as db:
db.execute(sql, (start_date, end_date))
total_articles = db.fetchall()
return total_articles
def add_journalist(self, first_name, last_name, email):
sql = 'Insert query to add a new user to the CMS'
with DatabaseHandler() as db:
db.execute(sql, (first_name, last_name, email))
Using @contextmanager
The @contextmanager decorator converts your generator function into a context manager so that it can be used with the withas control flow.
from contextlib import contextmanager
@contextmanager
def foobar():
print("Before")
yield {}
print("After")
with foobar() as fb:
print(fb)
"""
---OUTPUT---
Before
{}
After
"""
Yield is used to split the function into two halves. Everything before yield will be executed at the beginning of the with block and evrerything after yield will be called at the end of the block.
Even though you can’t see the __enter__ and __exit__ explicitly here, under the hood, they are being called.
__enter__before theyield__exit__before theyield
To learn more about how they are implemented, I did peek into some of the module's code.
Peeking into contextlib code
Going through the source code of contextlib I did find the __enter__ method and __exit__ method implementation.
def __enter__(self):
# do not keep args and kwds alive unnecessarily
# they are only needed for recreation, which is not possible anymore
del self.args, self.kwds, self.func
try:
return next(self.gen)
except StopIteration:
raise RuntimeError("generator didn't yield") from None
Since we are dealing with generator function here, the __enter__ method's main job is to next or yield the value of the gen object and return the yielded value so that it can be assigned to the variable after as in the withas block.
def __exit__(self, type, value, traceback):
if type is None:
try:
next(self.gen)
except StopIteration:
return False
else:
raise RuntimeError("generator didn't stop")
else:
if value is None:
value = type()
try:
self.gen.throw(type, value, traceback)
except StopIteration as exc:
return exc is not value
except RuntimeError as exc:
if exc is value:
return False
if type is StopIteration and exc.__cause__ is value:
return False
raise
except:
if sys.exc_info()[1] is value:
return False
raise
raise RuntimeError("generator didn't stop after throw()")
The exit method checks for an exception, if anything is passed to the type besides NONE, self.gen_throw is called causing the exception to be raised in the yield line. Otherwise next(gen) is called resuming the generator function body after the yield.
To make the code a bit readable on the screen I did remove the comments from the source code, feel free to check the source code on Github. I was amazed at how well commented the code was. Kudos to the devs for keeping the codebase so well maintained.
~ Real World Example with @contexmanager
@contextmanager
def database_handler():
host = 'localhost'
user = 'dev'
password = 'dev@123'
db = 'foobar'
port = '5432'
connection = None
cursor = None
connection = psycopg2.connect(
user=self.user,
password=self.password,
host=self.host,
port=self.port,
database=self.db
)
cursor = self.connection.cursor()
yield cursor
cursor.close()
connection.close()
The code snippet would work on a happy flow of the program but for any exception, this program is seriously flawed. You might remember the exit method in the contextlib source code, it raises exceptions if any in the yield line. Here we don’t seem to have any error handling in the yield line.
So during an exception generator function will abort without closing the resource properly. Failing to do the single job that it was meant for.
Tackling exceptions when using @contextmanager
@contextmanager
def database_handler():
try:
host = 'localhost'
user = 'dev'
password = 'dev@123'
db = 'foobar'
port = '5432'
connection = None
cursor = None
connection = psycopg2.connect(
user=self.user,
password=self.password,
host=self.host,
port=self.port,
database=self.db
)
cursor = self.connection.cursor()
yield cursor
except Exceptions as e:
logging.errror("Error connecting to the database", exc_info=True)
finally:
cursor.close()
connection.close()
By default __exit__ suppresses the exception. So exceptions should be reraised in the decorator function.
Having a try/finally (or a with block) around the yield is an unavoidable price of using @contextmanager, because you never know what the users of your context manager are going to do inside their with block. - Luciano Ramalho
PyRandom
The blog post is a follow up to the talk I gave at PythonPune about Context Managers. While researching for the talk I came across some interesting things hence I am adding this to the section PyRandom.
localcontextis great making high precision calculations in a particularwithasblock as when the flow exits the block, the precision is automatically restored.from decimal import localcontext with localcontext() as arthmetic: # Sets the current decimal precision to 42 places arthmetic.prec = 42 # Do some high_precision_arithmetic here # Automatically restores precision to the previous context arithmetic_calculation()When using multiple context managers the control flow seems to work in a LIFO, Last In First Out manner. The
__enter__method that is called last will have it’s__exit__method called first.import contextlib @contextlib.contextmanager def make_context(name): print ('entering:', name) yield name print ('exiting :', name) with make_context('A') as A, make_context('B') as B, make_context('C') as C: print ('inside with statement:', A, B, C) """ ---OUTPUT--- entering: A entering: B entering: C inside with statement: A B C exiting : C exiting : B exiting : A """Context managers created with
@contextmanagerare for single use.from contextlib import contextmanager @contextmanager def foobar(): print("Before") yield print("After") >>> foo = foobar() >>> with foo: pass Before After >>> with foo: pass Traceback (most recent call last): RuntimeError: generator didn't yieldContext Managers that can be nested with the
withcontrol flow and can use the same instance of a context manager in a nested flow are called Reetrant Context Managers.>>> from contextlib import redirect_stdout >>> from io import StringIO >>> stream = StringIO() >>> write_to_stream = redirect_stdout(stream) >>> with write_to_stream: ... print("This is written to the stream rather than stdout") ... with write_to_stream: ... print("This is also written to the stream") ... >>> print("This is written directly to stdout") This is written directly to stdout >>> print(stream.getvalue()) This is written to the stream rather than stdout This is also written to the streamContext Managers that cannot have a nested with control flow but their single instance can be used multiple times are called Reusable Context Managers
from threading import Lock lock = threading.Lock() # Lock gets acquired with lock: # Do some work # Lock gets released # Lock gets acquired with lock: # Do some work # Lock gets released # Lock gets acquired with lock: # Do some work with lock: # OOPS! Deadlock # Try to acquire an already acquired lock
Let's give credit where it's due. The talk and blog post both were created from the notes that I took two years back while reading Fluent Python by Luciano Ramalho.
I highly recommend that book if you want to upskill your Python Knowledge.
Happy Coding!
If you enjoy my content please consider supporting it.
Write a comment ...