import tempfile
import shutil


class TempDir(object):
    def __init__(self, delete=True, **kwargs):
        self.tempdir = tempfile.mkdtemp(**kwargs)
        self.delete = delete

    def __enter__(self):
        return self.tempdir

    def __exit__(self, type, value, traceback):
        if self.delete:
            shutil.rmtree(self.tempdir)