codekingpro/portable-devtools

1import contextlib2import os3import shutil4import stat5import tempfile6import time7from collections.abc import Callable, Generator8from functools import partial9from pathlib import Path10 11import yaml12from filelock import BaseFileLock, FileLock, SoftFileLock, Timeout13 14from .. import constants15from . import logging16 17 18logger = logging.get_logger(__name__)19 20# Wrap `yaml.dump` to set `allow_unicode=True` by default.21#22# Example:23# ```py24# >>> yaml.dump({"emoji": "👀", "some unicode": "日本か"})25# 'emoji: "\\U0001F440"\nsome unicode: "\\u65E5\\u672C\\u304B"\n'26#27# >>> yaml_dump({"emoji": "👀", "some unicode": "日本か"})28# 'emoji: "👀"\nsome unicode: "日本か"\n'29# ```30yaml_dump: Callable[..., str] = partial(yaml.dump, stream=None, allow_unicode=True)  # type: ignore31 32 33@contextlib.contextmanager34def SoftTemporaryDirectory(35    suffix: str | None = None,36    prefix: str | None = None,37    dir: Path | str | None = None,38    **kwargs,39) -> Generator[Path, None, None]:40    """41    Context manager to create a temporary directory and safely delete it.42 43    If tmp directory cannot be deleted normally, we set the WRITE permission and retry.44    If cleanup still fails, we give up but don't raise an exception. This is equivalent45    to  `tempfile.TemporaryDirectory(..., ignore_cleanup_errors=True)` introduced in46    Python 3.10.47 48    See https://www.scivision.dev/python-tempfile-permission-error-windows/.49    """50    tmpdir = tempfile.TemporaryDirectory(prefix=prefix, suffix=suffix, dir=dir, **kwargs)51    yield Path(tmpdir.name).resolve()52 53    try:54        # First once with normal cleanup55        shutil.rmtree(tmpdir.name)56    except Exception:57        # If failed, try to set write permission and retry58        try:59            shutil.rmtree(tmpdir.name, onerror=_set_write_permission_and_retry)60        except Exception:61            pass62 63    # And finally, cleanup the tmpdir.64    # If it fails again, give up but do not throw error65    try:66        tmpdir.cleanup()67    except Exception:68        pass69 70 71def _set_write_permission_and_retry(func, path, excinfo):72    os.chmod(path, stat.S_IWRITE)73    func(path)74 75 76@contextlib.contextmanager77def WeakFileLock(lock_file: str | Path, *, timeout: float | None = None) -> Generator[BaseFileLock, None, None]:78    """A filelock with some custom logic.79 80    This filelock is weaker than the default filelock in that:81    1. It won't raise an exception if release fails.82    2. It will default to a SoftFileLock if the filesystem does not support flock.83    3. Lock files are created with mode 0o664 (group-writable) instead of the default 0o644.84       This allows multiple users sharing a cache directory to wait for locks.85 86    An INFO log message is emitted every 10 seconds if the lock is not acquired immediately.87    If a timeout is provided, a `filelock.Timeout` exception is raised if the lock is not acquired within the timeout.88    """89    log_interval = constants.FILELOCK_LOG_EVERY_SECONDS90    lock = FileLock(lock_file, timeout=log_interval, mode=0o664)91    start_time = time.time()92 93    while True:94        elapsed_time = time.time() - start_time95        if timeout is not None and elapsed_time >= timeout:96            raise Timeout(str(lock_file))97 98        try:99            lock.acquire(timeout=min(log_interval, timeout - elapsed_time) if timeout else log_interval)100        except Timeout:101            logger.info(102                f"Still waiting to acquire lock on {lock_file} (elapsed: {time.time() - start_time:.1f} seconds)"103            )104        except NotImplementedError as e:105            if "use SoftFileLock instead" in str(e):106                logger.warning(107                    "FileSystem does not appear to support flock. Falling back to SoftFileLock for %s", lock_file108                )109                lock = SoftFileLock(lock_file, timeout=log_interval)110                continue111        else:112            break113 114    try:115        yield lock116    finally:117        try:118            lock.release()119        except OSError:120            try:121                Path(lock_file).unlink()122            except OSError:123                pass124