diff --git a/tests/data/test_make_nifti.py b/tests/data/test_make_nifti.py index f604eab76f..5ff4bac096 100644 --- a/tests/data/test_make_nifti.py +++ b/tests/data/test_make_nifti.py @@ -11,7 +11,11 @@ from __future__ import annotations +import atexit import os +import shutil +import subprocess +import sys import tempfile import unittest @@ -26,8 +30,10 @@ _, has_nib = optional_import("nibabel") TESTS = [] +caller_owned_dir = tempfile.mkdtemp() +atexit.register(shutil.rmtree, caller_owned_dir, ignore_errors=True) for affine in (None, np.eye(4), torch.eye(4)): - for dir in (None, tempfile.mkdtemp()): + for dir in (None, caller_owned_dir): for fname in (None, "fname"): TESTS.append([{"affine": affine, "dir": dir, "fname": fname}]) @@ -40,6 +46,21 @@ def test_make_nifti(self, params): created_file = make_nifti_image(im, verbose=True, **params) self.assertTrue(os.path.isfile(created_file)) + def test_temp_dir_removed_at_exit(self): + script = ( + "from monai.data.synthetic import create_test_image_2d;" + "from tests.test_utils import make_nifti_image;" + "print(make_nifti_image(create_test_image_2d(100, 88)[0]))" + ) + created_file = subprocess.check_output([sys.executable, "-c", script], text=True).strip() + self.assertFalse(os.path.exists(os.path.dirname(created_file))) + + def test_caller_owned_dir_kept(self): + im, _ = create_test_image_2d(100, 88) + with tempfile.TemporaryDirectory() as caller_dir: + make_nifti_image(im, dir=caller_dir) + self.assertTrue(os.path.isdir(caller_dir)) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_utils.py b/tests/test_utils.py index 5e21e48068..c4a3953d13 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -12,6 +12,7 @@ from __future__ import annotations import argparse +import atexit import copy import datetime import functools @@ -20,6 +21,7 @@ import operator import os import queue +import shutil import ssl import subprocess import sys @@ -391,7 +393,8 @@ def make_nifti_image( ): """ Create a temporary nifti image on the disk and return the image name. - User is responsible for deleting the temporary file when done with it. + If `dir` is not given, a temporary directory is created to hold the image and removed at + interpreter exit. If `dir` is given, the caller owns it. """ if isinstance(array, torch.Tensor): array, *_ = convert_data_type(array, np.ndarray) @@ -404,6 +407,7 @@ def make_nifti_image( # if dir not given, create random. Else, make sure it exists. if dir is None: dir = tempfile.mkdtemp() + atexit.register(shutil.rmtree, dir, ignore_errors=True) else: os.makedirs(dir, exist_ok=True)