TroyGrosfield.com TroyGrosfield.com

Headline

Python Context Manager and the With Statement

Author
by Troy Grosfield
Date
October 11th, 2010
Category
Developer
Story

Context managers are a great way to “loan” out objects. In conjunction with python’s “with” statement, this can be quite the powerful design pattern. You will see ways to automatically clean up entire objects or parts of objects once they are no longer needed. Below we create an object to be used temporarily, “loan” it out, then delete it or alter the object’s state in some way once we are finished with it.

Sample Code

from contextlib import contextmanager
import random

class SomeObject():
    """ Some Random class. """
    def __init__(self):
        self.id = random.randint(1000,9999)
        self.big_list = [i*1000 for i in range(1000)]
        print 'Initialize SomeObject. ID: {0}'.format(self.id)

    def __del__(self):
        print 'Calling destructor for ID: {0}'.format(self.id)
        self.big_list = None

@contextmanager
def create_some_object():
    """ Simple method to create and yield a new object. """
    so = SomeObject()

    try:
        print 'yield SomeObject. ID: {0}'.format(so.id)
        print so.big_list  # prints really big list
        yield so
    finally:
        print so.big_list  # prints really big list
        print 'Clean up SomeObject. ID: {0}'.format(so.id)
        so.__del__()
        print so.big_list  # prints None since the destructor was 
                              # called and the memory was cleaned up

# Now lets do something with the objects.
with create_some_object() as obj:
    # do something
    print 'Do something with SomeObject ID: {0}'.format(obj.id)

Output

Initialize SomeObject. ID: 9750
yield SomeObject. ID: 9750
{PRINTS BIG_LIST.  Omitted for purpose of saving space.}
Do something with SomeObject ID: 9750
{PRINTS BIG_LIST.  Omitted for purpose of saving space.}
Clean up SomeObject. ID: 9750
Calling destructor for ID: 9750
None
Calling destructor for ID: 9750  # Also called when the program exits

You’ll see for the demo purposes, I create SomeObject through the create_some_object() method. When the object is initialized, randomly generate a big_list that I only want to keep around for a short time. Then once I’m finished with the big_list portion of the object, I set it to None to free up the memory.

So how can you benefit from this design pattern? One practical use for context managers are for testing purposes. When you run your tests there will often be times where objects are created and you want your tests to clean up after themselves. With the method above, you can wrap your objects with the try finally logic and delete your object in the finally block once the test is finished with the object. That way when you run your tests your random test objects will be deleted from your database.

Sample Flow

def create_person():
    p = Person()
    p.id = 1

    try:
        yield p
    finally:
        p.delete()

with create_person() as p:
    self.assertEquals(1, p.id)

Assuming you’ve implemented a delete() method for your Person() object, the object will get deleted from your database when it’s no longer being used. This pattern will ensure cleanup of your object created during testing.

Note: when working with context managers be aware of memory usage. This would not be a good design pattern for load testing where you want to create a lot of objects and them all cleaned up. Your script will slow down immensely the more objects are created because they are all kept in memory and “loaned” out until they are no longer needed (you will likely need a very large number of items before this starts to slow down).

There are many different ways the functionality above can assist you. This is really just a simple introduction to what you can accomplish with python’s context managers and “with” statement.

Tags
Comments
No Comments »

No Comments Yet

Leave a reply