在Python装饰器中创建参数以供装饰函数使用

Question

I have several functions that use tempfile.mkstemp() to create and use temporary files that stay on disk after the function has been called. Each of them repeats the same scaffolding code to clean up the file descriptor when the function exits. For example:

import tempfile
import zipfile
import os

def write_zip(test_name: str) -> str:
    """
    Given a test name, returns an archive (ZIP) of the files in that test.
    """

    try:
        fd, zip_path = tempfile.mkstemp(suffix='.zip')

        with zipfile.ZipFile(zip_path, 'w') as zf:
            for fpath in _list_files(test_name):  # _list_files() not shown
                zf.write(fpath, arcname=os.path.basename(fpath))
    finally:
        try:
            os.close(locals()['fd'])
        except KeyError:
            pass

    return zip_path

I'd like to pull this try-finally scaffolding out to a decorator. I tried writing one, but I'm a bit lost and it isn't working:

def _persistent_temp_file(func):
    """
    Decorator for a function that uses a temporary file that should persist on disk after the function has exited.
    Closes the file descriptor in a try-finally block so the function doesn't have to. 
    """

    def scaffolding(suffix=None, prefix=None, dir=None, text=False):
        try:
            fd, temp_file = tempfile.mkstemp(suffix=suffix, prefix=prefix, dir=dir, text=text)
            func(temp_file)

        finally:
            try:
                os.close(locals()['fd'])
            except KeyError:
                pass

    return scaffolding


@_persistent_temp_file
def write_zip(test_name: str) -> str:
    with zipfile.ZipFile(zip_path, 'w') as zf:
        for fpath in _list_files(test_name):  # _list_files() not shown
            zf.write(fpath, arcname=os.path.basename(fpath))

    return zip_path

I'm unsure of how to pass the zip_file parameter to this function (or any other decorated function). I'm not sure I know how to pass in the parameters needed by mkstemp(), either. (ie How do I specify the file suffix when calling the decorated function?)

Answer 1

Rather than use a decorator, you'd want to create a context manager . A context manager is told when a block of code (the context) exits, so you can clean up after yourself at that point.

Ironically, there is a decorator that makes it easy to write a context manager, called @contextlib.contextmanager :

from contextlib import contextmanager

@contextmanager
def _persistent_temp_file(suffix=None, prefix=None, dir=None, text=False):
    fd = None
    try:
        fd, temp_file = tempfile.mkstemp(suffix=suffix, prefix=prefix, dir=dir, text=text)
        yield temp_file
    finally:
        if fd is not None:
            os.close(fd)

Note the yield temp_file line; that's the point at which this function is paused, and the temp_file value is then returned from the resulting context manager __enter__ method, and becomes available to the as clause.

Then use this in a with statement:

def write_zip(test_name: str) -> str:
    with _persistent_temp_file() as zip_path:
        with zipfile.ZipFile(zip_path, 'w') as zf:
            for fpath in _list_files(test_name):  # _list_files() not shown
                zf.write(fpath, arcname=os.path.basename(fpath))
    return zip_path

Not that you need to re-invent the persistent-temp-file wheel here. You could just use the tempfile.NamedTempFile() object here instead:

from tempfile import NamedTempFile

def write_zip(test_name: str) -> str:
    with NamedTempFile(delete=False) as temp_file:
        with zipfile.ZipFile(temp_file, 'w') as zf:
            for fpath in _list_files(test_name):  # _list_files() not shown
                zf.write(fpath, arcname=os.path.basename(fpath))
    return temp_file.name