Adding Temporary Directory Routine

monai/apps/__init__.py

@@ -13,4 +13,13 @@
 
 from .datasets import CrossValidation, DecathlonDataset, MedNISTDataset, TciaDataset
 from .mmars import MODEL_DESC, RemoteMMARKeys, download_mmar, get_model_spec, load_from_mmar
-from .utils import SUPPORTED_HASH_TYPES, check_hash, download_and_extract, download_url, extractall, get_logger, logger
+from .utils import (
+    SUPPORTED_HASH_TYPES,
+    check_hash,
+    create_temp_dir,
+    download_and_extract,
+    download_url,
+    extractall,
+    get_logger,
+    logger,
+)

monai/apps/utils.py

@@ -11,6 +11,7 @@
 
 from __future__ import annotations
 
+import atexit
 import hashlib
 import json
 import logging
@@ -29,7 +30,7 @@
 from urllib.request import urlopen, urlretrieve
 
 from monai.config.type_definitions import PathLike
-from monai.utils import look_up_option, min_version, optional_import
+from monai.utils import MONAIEnvVars, look_up_option, min_version, optional_import
 
 requests, has_requests = optional_import("requests")
 gdown, has_gdown = optional_import("gdown", "4.7.3")
@@ -42,7 +43,15 @@
 else:
     tqdm, has_tqdm = optional_import("tqdm", "4.47.0", min_version, "tqdm")
 
-__all__ = ["check_hash", "download_url", "extractall", "download_and_extract", "get_logger", "SUPPORTED_HASH_TYPES"]
+__all__ = [
+    "check_hash",
+    "download_url",
+    "extractall",
+    "download_and_extract",
+    "get_logger",
+    "SUPPORTED_HASH_TYPES",
+    "create_temp_dir",
+]
 
 DEFAULT_FMT = "%(asctime)s - %(levelname)s - %(message)s"
 SUPPORTED_HASH_TYPES = {"md5": hashlib.md5, "sha1": hashlib.sha1, "sha256": hashlib.sha256, "sha512": hashlib.sha512}
@@ -422,3 +431,34 @@ def download_and_extract(
         filename = filepath or Path(tmp_dir, get_filename_from_url(url)).resolve()
         download_url(url=url, filepath=filename, hash_val=hash_val, hash_type=hash_type, progress=progress)
         extractall(filepath=filename, output_dir=output_dir, file_type=file_type, has_base=has_base)
+
+
+def create_temp_dir(directory: str | None = None, delete_on_finalise: bool = False) -> str:
+    """
+    Creates or uses an existing temporary directory. If `directory` is given, this is used as the path to a directory
+    which is created if it doesn't exist already. If `directory` is None, the value of the environment variable
+    `MONAI_DATA_DIRECTORY` is used instead, if this is not present then a random temporary directory is chosen. If
+    `delete_on_finalise` is True, or a random temp directory was created, the directory and its contents will be deleted
+    when the process exits.
+
+    Args:
+        directory: path to desired temporary directory, or None to use MONAI_DATA_DIRECTORY or choose a random one.
+        delete_on_finalise: if True, the directory and its contents will be deleted when the process exits.
+
+    Returns:
+        The path to the existing or new temporary directory, if `directory` is given this will be returned, otherwise it
+        will be the value of MONAI_DATA_DIRECTORY if present or the chosen random directory. This directory will exist.
+
+    """
+    if directory is None:
+        directory = MONAIEnvVars.data_dir()
+        if directory is None:
+            directory = tempfile.mkdtemp()
+            delete_on_finalise = True
+
+    os.makedirs(directory, exist_ok=True)
+
+    if delete_on_finalise:
+        atexit.register(shutil.rmtree, directory, ignore_errors=True)
+
+    return directory

tests/apps/test_create_temp_dir.py

@@ -0,0 +1,84 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import os
+import tempfile
+import unittest
+from unittest.mock import patch
+
+from monai.apps import create_temp_dir
+from monai.utils import MONAIEnvVars
+
+MONAI_DATA_DIRECTORY = "MONAI_DATA_DIRECTORY"
+
+
+class TestCreateTempDir(unittest.TestCase):
+    def test_basic_use(self):
+        """Test basic usage which should create a new random temporary directory."""
+        try:
+            data_dir = os.environ.pop(MONAI_DATA_DIRECTORY, None)  # ignore the environment variable if present
+
+            test_dir = create_temp_dir()
+
+            self.assertTrue(os.path.isdir(test_dir))
+
+        finally:
+            if data_dir is not None:
+                os.environ[MONAI_DATA_DIRECTORY] = data_dir
+
+    def test_data_dir(self):
+        """Test using a mocked MONAI_DATA_DIRECTORY, which should be returned by the function."""
+        with patch("monai.utils.MONAIEnvVars.data_dir") as data_dir, tempfile.TemporaryDirectory() as fake_data_dir:
+            data_dir.return_value = fake_data_dir
+
+            self.assertEqual(fake_data_dir, MONAIEnvVars.data_dir())
+
+            test_dir = create_temp_dir()
+
+            self.assertTrue(os.path.isdir(test_dir))
+            self.assertEqual(test_dir, fake_data_dir)
+
+    def test_given_dir(self):
+        """Test giving a directory to the function, ensuring it creates the directory."""
+        with tempfile.TemporaryDirectory() as temp_dir:
+            selected_dir = f"{temp_dir}/test_inner_dir"
+            test_dir = create_temp_dir(selected_dir)
+
+            self.assertTrue(os.path.isdir(selected_dir))
+            self.assertEqual(test_dir, selected_dir)
+
+    def test_finalisation(self):
+        """Test the temporary directory is deleted by finalisation using a subprocess."""
+        self.finaliser = None
+
+        def _register(func, /, *args, **kwargs):
+            self.finaliser = (func, args, kwargs)
+
+        with patch("atexit.register", new=_register), tempfile.TemporaryDirectory() as temp_dir:
+            selected_dir = f"{temp_dir}/test_inner_dir"
+            test_dir = create_temp_dir(selected_dir, True)
+
+            self.assertTrue(os.path.isdir(selected_dir))
+            self.assertIsNotNone(self.finaliser)
+
+            with open(test_dir + "/test_file", "w") as o:
+                o.write("Test file data")
+
+            func, args, kwargs = self.finaliser
+            func(*args, **kwargs)
+
+            self.assertFalse(os.path.exists(selected_dir))
+
+
+if __name__ == "__main__":
+    unittest.main()