Question

I’m trying to create an object that plays nicely with the with keyword in Python. I understand that you have to create __enter__ and __exit__ methods, but I’m not quite sure how manipulate the object. As a concrete example I wrote a class that creates local space to work in, and cleans up on exit:

import tempfile, os, shutil
class temp_workspace(object):

    def __enter__(self):
        self.local_dir = os.getcwd()
        self.temp_dir  = tempfile.mkdtemp()
        os.chdir(self.temp_dir)

    def __exit__(self, exc_type, exc_value, traceback):
        os.chdir(self.local_dir)
        shutil.rmtree(self.temp_dir)

    def __repr__(self):
        return self.temp_dir

This works just fine, but when I try to print the local directory name:

with temp_workspace() as T:
    print "Temp directory name is ", T

It shows up as None and __repr__ isn’t even called! It’s really confusing since T is also NoneType. What am I doing wrong?

Answer 1

You’re not returning the object from __enter__ as specified by the context manager protocol. Add return self to the end of your __enter__ method.