codekingpro/portable-devtools

1"""holds locking functionality that works across processes."""2 3from __future__ import annotations4 5import logging6import os7from abc import ABC, abstractmethod8from contextlib import contextmanager, suppress9from pathlib import Path10from threading import Lock, RLock11 12from filelock import FileLock, Timeout13 14 15class _CountedFileLock(FileLock):16    def __init__(self, lock_file) -> None:17        parent = os.path.dirname(lock_file)18        if not os.path.isdir(parent):19            with suppress(OSError):20                os.makedirs(parent)21 22        super().__init__(lock_file)23        self.count = 024        self.thread_safe = RLock()25 26    def acquire(self, timeout=None, poll_interval=0.05):27        if not self.thread_safe.acquire(timeout=-1 if timeout is None else timeout):28            raise Timeout(self.lock_file)29        if self.count == 0:30            super().acquire(timeout, poll_interval)31        self.count += 132 33    def release(self, force=False):  # noqa: FBT00234        with self.thread_safe:35            if self.count > 0:36                self.thread_safe.release()37            if self.count == 1:38                super().release(force=force)39            self.count = max(self.count - 1, 0)40 41 42_lock_store = {}43_store_lock = Lock()44 45 46class PathLockBase(ABC):47    def __init__(self, folder) -> None:48        path = Path(folder)49        self.path = path.resolve() if path.exists() else path50 51    def __repr__(self) -> str:52        return f"{self.__class__.__name__}({self.path})"53 54    def __truediv__(self, other):55        return type(self)(self.path / other)56 57    @abstractmethod58    def __enter__(self):59        raise NotImplementedError60 61    @abstractmethod62    def __exit__(self, exc_type, exc_val, exc_tb):63        raise NotImplementedError64 65    @abstractmethod66    @contextmanager67    def lock_for_key(self, name, no_block=False):  # noqa: FBT00268        raise NotImplementedError69 70    @abstractmethod71    @contextmanager72    def non_reentrant_lock_for_key(self, name):73        raise NotImplementedError74 75 76class ReentrantFileLock(PathLockBase):77    def __init__(self, folder) -> None:78        super().__init__(folder)79        self._lock = None80 81    def _create_lock(self, name=""):82        lock_file = str(self.path / f"{name}.lock")83        with _store_lock:84            if lock_file not in _lock_store:85                _lock_store[lock_file] = _CountedFileLock(lock_file)86            return _lock_store[lock_file]87 88    @staticmethod89    def _del_lock(lock):90        if lock is not None:91            with _store_lock, lock.thread_safe:92                if lock.count == 0:93                    _lock_store.pop(lock.lock_file, None)94 95    def __del__(self) -> None:96        self._del_lock(self._lock)97 98    def __enter__(self):99        self._lock = self._create_lock()100        self._lock_file(self._lock)101 102    def __exit__(self, exc_type, exc_val, exc_tb):103        self._release(self._lock)104        self._del_lock(self._lock)105        self._lock = None106 107    def _lock_file(self, lock, no_block=False):  # noqa: FBT002108        # multiple processes might be trying to get a first lock... so we cannot check if this directory exist without109        # a lock, but that lock might then become expensive, and it's not clear where that lock should live.110        # Instead here we just ignore if we fail to create the directory.111        with suppress(OSError):112            os.makedirs(str(self.path))113 114        try:115            lock.acquire(0.0001)116        except Timeout:117            if no_block:118                raise119            logging.debug("lock file %s present, will block until released", lock.lock_file)120            lock.release()  # release the acquire try from above121            lock.acquire()122 123    @staticmethod124    def _release(lock):125        lock.release()126 127    @contextmanager128    def lock_for_key(self, name, no_block=False):  # noqa: FBT002129        lock = self._create_lock(name)130        try:131            try:132                self._lock_file(lock, no_block)133                yield134            finally:135                self._release(lock)136        finally:137            self._del_lock(lock)138            lock = None139 140    @contextmanager141    def non_reentrant_lock_for_key(self, name):142        with _CountedFileLock(str(self.path / f"{name}.lock")):143            yield144 145 146class NoOpFileLock(PathLockBase):147    def __enter__(self):148        raise NotImplementedError149 150    def __exit__(self, exc_type, exc_val, exc_tb):151        raise NotImplementedError152 153    @contextmanager154    def lock_for_key(self, name, no_block=False):  # noqa: ARG002, FBT002155        yield156 157    @contextmanager158    def non_reentrant_lock_for_key(self, name):  # noqa: ARG002159        yield160 161 162__all__ = [163    "NoOpFileLock",164    "ReentrantFileLock",165    "Timeout",166]167