我有几个函数使用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
我想把这个try-finally脚手架拉到装饰器上。我试过写一个,但我有点迷失了,它不起作用:
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()所需的参数。 (即如何在调用修饰函数时指定文件后缀?)
答案 0 :(得分:1)
您不想使用装饰器,而是想要创建context manager。当一个代码块(上下文)退出时,会告诉上下文管理器,因此您可以在此之后自行清理。
具有讽刺意味的是,有一个装饰器可以很容易地编写一个名为@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