How to define a Python class that I can use with statement c?

Question

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

Now I want to use StringIO with a with statement:

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

But Python complains that the StringIO type StringIO not have a __exit__ method. After subclassing StringIO :

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

Now I get an exception due to the lack of the __enter__ method. How to define the __enter__ method? What does Python expect from a class that can be used with a with statement?

Answer 1

You need to write a context manager . If you do not want to write the entire protocol, simplify it using the contextlib.contextmanager decorator.

Answer 2

Just to give an example where this would be redundant, use contextlib, which is more for use with generators:

Here is a good and useful example of a context manager. I found it on David Basley's blog:

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

and modified it a bit to use the best timer on the platform, the trick 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 this gives you a context manager that you can use as follows:

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

and you will get a printout. I liked this example and would like to share it.

Answer 3

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