在Python裝飾器中創建參數以供裝飾函數使用

Question

我有幾個使用tempfile.mkstemp（）的函數來創建和使用在函數被調用后保留在磁盤上的臨時文件。 當函數退出時，它們每個都重復相同的腳手架代碼以清理文件描述符。 例如：

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

我想把這個最終嘗試的腳手架拉到一個裝飾器上。 我嘗試寫一個，但是我有點迷茫，它不起作用：

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

我不確定如何將zip_file參數傳遞給此函數（或任何其他裝飾函數）。 我不確定我也不知道如何傳遞mkstemp（）所需的參數。 （即調用裝飾函數時如何指定文件后綴？）

Answer 1

您想要創建一個上下文管理器，而不是使用裝飾器 。 當代碼塊（上下文）退出時，將通知上下文管理器，因此您可以在此之后自行清理。

具有諷刺意味的是，有一個裝飾器可以輕松地編寫上下文管理器，稱為@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)

注意yield temp_file行； 這是在其中該功能被暫停的點，並且temp_file值然后從所得的上下文管理器返回__enter__方法，並且變得可用於as子句。

然后在with語句中使用它：

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

並不是說您需要在這里重新發明持久性臨時文件輪。 您可以只在這里使用tempfile.NamedTempFile()對象：

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