decorators.py 文件源码

def make_tempdir(prefix=None, suffix=None):
    """Decorator that adds a last argument that is the path to a temporary
    directory that gets deleted after the function has finished.

    Usage::

        @make_tempdir()
        def some_function(arg1, arg2, tempdir, kwargs1='one'):
            assert os.path.isdir(tempdir)
            ...
    """

    def decorator(func):

        @wraps(func)
        def inner(*args, **kwargs):
            with TemporaryDirectory(prefix=prefix, suffix=suffix) as f:
                args = args + (f,)
                return func(*args, **kwargs)

        return inner

    return decorator