diff --git a/src/utils/file_lock.py b/src/utils/file_lock.py index f2c1aa13..0cbc0203 100644 --- a/src/utils/file_lock.py +++ b/src/utils/file_lock.py @@ -7,34 +7,36 @@ from pathlib import Path import portalocker +def _get_portalocker_lock( + path: Path, flags: int, timeout: Optional[float] = None +) -> portalocker.Lock: + """Helper to prepare path and create a portalocker.Lock instance.""" + path.parent.mkdir(parents=True, exist_ok=True) + path.touch(exist_ok=True) + + kwargs = { + "mode": "r+b", + "flags": flags, + } + if timeout is not None: + kwargs["timeout"] = timeout + + return portalocker.Lock(str(path), **kwargs) + + @contextmanager def exclusive_lock( path: Union[str, PathLike[str], Path], timeout: Optional[float] = None ) -> Generator[None, None, None]: """Context manager that locks *path* exclusively. - The function opens the file in binary append mode and obtains an + The function opens the file in binary read/write mode and obtains an exclusive lock using ``portalocker``. If ``timeout`` is provided, acquiring the lock will wait for at most that many seconds before raising ``portalocker.exceptions.LockException``. """ path = Path(path) - path.parent.mkdir(parents=True, exist_ok=True) - path.touch(exist_ok=True) - flags = portalocker.LockFlags.EXCLUSIVE - if timeout is None: - lock = portalocker.Lock( - str(path), - mode="r+b", - flags=flags, - ) - else: - lock = portalocker.Lock( - str(path), - mode="r+b", - timeout=timeout, - flags=flags, - ) + lock = _get_portalocker_lock(path, portalocker.LockFlags.EXCLUSIVE, timeout) with lock as fh: yield fh @@ -51,21 +53,7 @@ def shared_lock( ``portalocker.exceptions.LockException``. """ path = Path(path) - path.parent.mkdir(parents=True, exist_ok=True) - path.touch(exist_ok=True) - if timeout is None: - lock = portalocker.Lock( - str(path), - mode="r+b", - flags=portalocker.LockFlags.SHARED, - ) - else: - lock = portalocker.Lock( - str(path), - mode="r+b", - timeout=timeout, - flags=portalocker.LockFlags.SHARED, - ) + lock = _get_portalocker_lock(path, portalocker.LockFlags.SHARED, timeout) with lock as fh: fh.seek(0) yield fh