Class TempDirectory

python/tvm/support/utils.py:39–178

Helper object to manage temp directory during testing. Automatically removes the directory when it went out of scope.

class TempDirectory:
    """Helper object to manage temp directory during testing.

    Automatically removes the directory when it went out of scope.
    """

    # When True, all TempDirectory are *NOT* deleted and instead live inside a predicable directory
    # tree.
    _KEEP_FOR_DEBUG = False

    # In debug mode, each tempdir is named after the sequence
    _NUM_TEMPDIR_CREATED = 0
    _NUM_TEMPDIR_CREATED_LOCK = threading.Lock()

    @classmethod
    def _increment_num_tempdir_created(cls):
        with cls._NUM_TEMPDIR_CREATED_LOCK:
            to_return = cls._NUM_TEMPDIR_CREATED
            cls._NUM_TEMPDIR_CREATED += 1

        return to_return

    _DEBUG_PARENT_DIR = None

    @classmethod
    def _get_debug_parent_dir(cls):
        if cls._DEBUG_PARENT_DIR is None:
            all_parents = f"{tempfile.gettempdir()}/tvm-debug-mode-tempdirs"
            if not os.path.isdir(all_parents):
                os.makedirs(all_parents)
            cls._DEBUG_PARENT_DIR = tempfile.mkdtemp(
                prefix=datetime.datetime.now().strftime("%Y-%m-%dT%H-%M-%S___"), dir=all_parents
            )
        return cls._DEBUG_PARENT_DIR

    TEMPDIRS = set()

    @classmethod
    def remove_tempdirs(cls):
        temp_dirs = getattr(cls, "TEMPDIRS", None)
        if temp_dirs is None:
            return

        for path in temp_dirs:
            shutil.rmtree(path, ignore_errors=True)

        cls.TEMPDIRS = None

    @classmethod
    @contextlib.contextmanager
    def set_keep_for_debug(cls, set_to=True):
        """Keep temporary directories past program exit for debugging."""
        old_keep_for_debug = cls._KEEP_FOR_DEBUG
        try:
            cls._KEEP_FOR_DEBUG = set_to
            yield
        finally:
            cls._KEEP_FOR_DEBUG = old_keep_for_debug