Create mkdtemp context manager

[David Wolever]

Something along the lines of:

class mkdtemp(object):
    def __init__(self, *args, **kwargs):
        self.should_chdir = kwargs.pop("chdir")
        self.original_dir = None
        self.dirname = tempfile.mkdtemp(*args, **kwargs)

    def __enter__(self):
        if self.should_chdir:
            self.original_dir = os.getcwd()
            os.chdir(self.dirname)
        return self.dirname

    def __exit__(self, *exc):
        if self.original_dir:
            os.chdir(self.original_dir)
        shutil.rmtree(self.dirname)

with mkdtemp("foo", chdir=True) as dirname:
    print "I'm in", dirname

[Cory Benfield]

I wondered about returning self from __enter__: it might be better to return self.dirname, really.

[David Wolever]

D'oh! Yes, that does make more sense.

[Andrey Petrov]

I love that this has been more than a year, but I'm still +1.