utils.py

from pathlib import Path
import tempfile
from contextlib import contextmanager
import os


def get_drive(path: str):
    path = Path(path).resolve()
    drive = path.drive
    root = path.root
    return drive + root


@contextmanager
def custom_drive_cache_dir(drive: str):
    drive = Path(drive)
    base_dir = Path(drive) / "tmp"
    if not base_dir.exists():
        os.makedirs(base_dir)
    print(f"Using {base_dir.resolve()} as cache dir")
    with tempfile.TemporaryDirectory(dir=base_dir) as tmp_dir:
        yield tmp_dir