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.