Safer directory creation

safe_mkdir.py

@contextlib.contextmanager
def safe_mkdir(path, dir=None):
    """ Create a directory in safe(r) way. Specified directory is created only when
    whole operations in `with` scoped is completed successfully. All the files
    that are created within the temporaly generated dir will be kept within.

    This may not work in some OS.
    """
    if dir is None:
        dir = os.path.dirname(path)
    if os.path.exists(path):
        raise OSError("Path '%s' already exists" % path)
    d = tempfile.mkdtemp(prefix=os.path.split(path)[-1], dir=dir)
    try:
        yield d
    except:
        shutil.rmtree(d)
        raise
    # Test again because shutil.move will not fail even if a directory exists
    if os.path.exists(path):
        raise OSError("Path '%s' already exists" % path)
    try:
        shutil.move(d, path)
    except:
        shutil.rmtree(d)
        raise


# A directory 'tmp' will be created with file 'tmp/bar.txt'
with safe_output_dir('tmp') as d:
    with open(os.path.join(d, 'bar.txt'), 'w') as fout:
        fout.write('foo')

# No directory will be created
with safe_output_dir('tmp') as d:
    1/0
    with open(os.path.join(d, 'bar.txt'), 'w') as fout:
        fout.write('foo')