Simple rollback on error with contextlib.ExitStack

rollback_with_exitstack.py

import os
import shutil
from contextlib import ExitStack


def create_folder(stack: ExitStack, path):
    try:
        os.mkdir(path)
    except FileExistsError:
        pass

    if not os.path.isdir(path):
        raise NotADirectoryError(repr(path))

    stack.callback(lambda: os.rmdir(path))


def copy_file(stack: ExitStack, src, dst):
    shutil.copyfile(src, dst)
    stack.callback(lambda: os.unlink(dst))


def main():
    with ExitStack() as stack:
        create_folder(stack, './foobar')
        copy_file(stack, './foo.json', './foobar/spam.json')
        copy_file(stack, './foo.json', './foobar/spam1.json')
        copy_file(stack, './foo.json', './foobar/spam2.json')
        copy_file(stack, './foo.json', './foobar/spam3.json')
        copy_file(stack, './foo.json', './foobar/spam4')
        create_folder(stack,'./foobar/spam4')  # Error!
        create_folder(stack, './foobar/spam5')
        stack.pop_all()  # on success, clear exit stack, so the files/folders stay.


if __name__ == '__main__':
    main()