tests: remove the temporary directory make_nifti_image creates

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()

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)