Python Context Managers and the “with” Statement

If you’re like me and you are relatively new to Python, you may have come across the term “context manager” while reading documentation, or researching a library that you are thinking about using. For a while, I glossed over this term to get what I needed done, but in the back of my mind I was still curious. Along the same lines, Python has this with keyword. What’s up with that thing?

In this post, I will go over what I have learned about context managers, the with keyword, and how they are related.

As it turns out, context managers and the with keyword go hand-in-hand!

I had read here and there that when you want to read data from a file, it is better to use the with keyword to do so because it automatically closes the file after you leave the scope of the “with statement.” This as opposed to having to remember to close the file when you’re done with it.

with open('filename.txt') as f:
    text = f.read()
    print(text)
# File is closed by this point

Yeah yeah, but still, what is a context manager?

Context managers are just Python classes that specify the __enter__ and __exit__ methods. Actually, as PEP 343 states:

This PEP adds a new statement "with" to the Python language to make
it possible to factor out standard uses of try/finally statements.

In this PEP, context managers provide __enter__() and __exit__()
methods that are invoked on entry to and exit from the body of the
with statement.

As lack-luster as that sounds, just specifying these two methods in a class makes that class a “context manager.” As stated, the __enter__ method is where you put the code that is going to run before you enter the scope of the with block. Similarly, the __exit__ method is where you put the code that you want to execute when you leave the scope of the with block.

In the example above, open() is opening the specified file, then we do what we need with the contents of the file, then we leave the with block and the file is closed without us having to do any extra work.

This is nice and all, but if you think about it, this concept can be applied to other scenarios where you want to have a reusable way to create a “state of the world” for code to run in, rather than typing the same thing over and over. An example of where context managers could be useful is in creating database connections, and then closing them when finished.

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

class SQLAlchemyDBConnection(object):
    """SQLAlchemy database connection"""
    def __init__(self, connection_string):
        self.connection_string = connection_string
        self.session = None
    def __enter__(self):
        engine = create_engine(self.connection_string)
        Session = sessionmaker()
        self.session = Session(bind=engine)
        return self
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.session.close()

Because we specified the __enter__ and __exit__ methods, we can use this class with the with keyword.

conn_str = 'mssql+pydobc://server_name/db_name?driver=SQL+Server'
with SQLAlchemyDBConnection(conn_str) as db:
    customer = db.session.query(Customer).filter_by(id=123).one()
    print(customer.name)

Alternatively, you can also write the with statement this way:

conn_str = 'mssql+pydobc://server_name/db_name?driver=SQL+Server'
db = SQLAlchemyDBConnection(conn_str)
with db:
    customer = db.session.query(Customer).filter_by(id=123).one()
    print(customer.name)

In these examples I’ve used a SQLAlchemy database connection, but you could create one for MongoDB, or some other database querying tool.

from pymongo import MongoClient

class MongoDBConnection(object):
    """MongoDB Connection"""
    def __init__(self, host='localhost', port='27017')
        self.host = host
        self.port = port
        self.connection = None
    def __enter__(self):
        self.connection = MongoClient(self.host, self.port)
        return self
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.connection.close()

We’d use this in a similar way.

mongo = MongoDBConnection()
with mongo:
    collection = mongo.connection.MyDbName.Customers
    customer = collection.find({'_id': 123})
    print(customer.get('name'))

But wait! There’s more!

Not only is this awesome for reusability, but it turns out that you can have your code be within the context of multiple context managers at the same time!

sql_conn_str = 'mssql+pydobc://server_name/db_name?driver=SQL+Server'
mongo = MongoDbConnection('10.9.8.7', 3000)
sql = SQLAlchemyDBConnection(sql_conn_str)
with mongo, sql:
    # get data from SQL
    customer = sql.session.query(Customer).first()
    # insert data into mongo
    collection = mongo.connection.DbName.CollectionName
    collection.insert({
        customer_id: customer.cust_id,
        name: customer.name})

You can “nest” as many context managers as you like using the with statement. In fact, rather than comma-separating them, you could actually nest them.

with MongoDbConnection() as mongo:
    with SQLAlchemyDBConnection(conn_str) as sql:
        # Do stuff

You can also comma-separate this longer version. I prefer to save them to variables first and comma-separate those variables. It’s really up to you.

Speaking of it’s-up-to-you, you don’t have to do this with database connections, nor am I suggesting that you should. I just thought it a good example of setting up a particular state for code to run in, and then doing away with it.

Also, as I mentioned towards the beginning, you might have come across the term “context manager” in documentation, but weren’t sure what it meant at the time. But hopefully, now you’ll have a better understanding of what’s going on.

Here is an example from the requests library.

For those that want to get deeper into the weeds, check out the docs on the with statement. There is also the contextlib library in the Python standard library that specifies a contextmanager decorator. This let’s you turn a function into a context manager, but I’ll let you read up on that.

While I have seen posts by others about context managers, I hadn’t come across a post about nesting them in with statements, outside of the documentation. I’ve not explored this very deeply, but it seems that interesting scenarios are possible where your code could be within the context of a database connection and a requests Session. That might turn out to be a ridiculous scenario, but hey, you could do it!

Anyway, hopefully someone finds this info useful. Thanks for reading!