How do I define a Python class that I can use with the with statement?

Question

I understand that StringIO acts like a file object, duck-typing what you would get from open('somefile.txt') .

Now I want to use StringIO with the with statement:

with StringIO('some string') as fh: # fh as in "file handle"
    data = [stuff from stuff in fh.read()]

But Python complains that type StringIO does not have an __exit__ method. After subclassing StringIO :

class MyStringIO(StringIO):
    def __exit__(self):
        self.close()

I now get an exception about not having an __enter__ method. How do I define the __enter__ method? What does Python expect from a class that can be used with the with statement?

Answer 1

You need to write a context manager . If you don't want to write the whole protocol, there's a simplified way around it using the contextlib.contextmanager decorator.

Answer 2

Just to give an example where it would be overkill do use contextlib, which is more for use with generators:

Here is a nice and useful example of a context manager. I found it in David Beazley's blog:

http://dabeaz.blogspot.de/2010/02/context-manager-for-timing-benchmarks.html

and just modified it a little bit to use the best timer per platform, a trick that I use from the timeit module:

# benchmark.py
import timeit
class benchmark(object):
    from timeit import default_timer as timer
    def __init__(self, name):
        self.name = name
    def __enter__(self):
        self.start = self.timer()
    def __exit__(self, ty, val, tb):
        end = self.timer()
        print("%s : %0.3f seconds" % (self.name, end-self.start))
        return False

Now that gives you a context manager that you can use like this:

with benchmark('algorithm 1'):
    ... (some computation)
with benchmark('other approach'):
    ... (some other computation)

and you get your output printed. I liked this example and wanted to share it.

Answer 3

you define the __enter__ method just as you did with __exit__ . these are the two methods needed for an object to be used with the with statement, see this for more info.