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