Python Returns Inside With

The Python with statement, introduced in 2.5, is awesome.  I just discovered, perhaps a little naively, how awesome it really is.

For those that don't know, the with statement allows you to use an object inside a new context.  You can't just use any only Python object, like a string or a dictionary.  Any object used in a with statement needs a context manager.  For more on using context managers (and the with statement in general), see Fredrik Lundh's guide.

The additional piece of with magic I've been playing with is returning data from inside a with statement.  For example, consider the following getcount() function that returns that line count of a file...

def getcount(fobj):
    with fobj as f:
        return len(f.readlines())

fobj = open('tmp.txt')
print getcount(fobj)
print fobj.closed

This example is quite simple — open a file object and print the number of lines.  The getcount() function uses the with statement to open and close the file.  But wait a minute, how was the file object closed if we returned from the function inside the with statement?  This is the beauty of the context manager — it is guaranteed to carry out all clean-up tasks no matter what.  This is why we're printing out fobj.closed — to verify that this did indeed work.

The same methodology can be used with lock context managers.  For example...

import time
from threading import Thread, Lock

class MyThread(Thread):

    lock = Lock()

    def __init__(self):
        super(MyThread, self).__init__()

    def set_name(self):
        with self.lock:
            self.name = 'MyName'
            time.sleep(2)

    def get_name(self):
        print 'Waiting for name...'
        with self.lock:
            return self.name

    def run(self):
        self.set_name()

tobj = MyThread()
tobj.start()
print tobj.get_name()

Here we're explicitly making set_name() take longer than it should — causing the lock object to block.  The return statement in get_name() is paused until the name is set, since it's inside the with.  Cool stuff.