zulip/zerver/lib/context_managers.py

51 lines
1.5 KiB
Python

"""
Context managers, i.e. things you can use with the 'with' statement.
"""
import fcntl
from collections.abc import Iterator
from contextlib import contextmanager
from typing import IO, Any
@contextmanager
def flock(lockfile: int | IO[Any], shared: bool = False) -> Iterator[None]:
"""Lock a file object using flock(2) for the duration of a 'with' statement.
If shared is True, use a LOCK_SH lock, otherwise LOCK_EX."""
fcntl.flock(lockfile, fcntl.LOCK_SH if shared else fcntl.LOCK_EX)
try:
yield
finally:
fcntl.flock(lockfile, fcntl.LOCK_UN)
@contextmanager
def lockfile(filename: str, shared: bool = False) -> Iterator[None]:
"""Lock a file using flock(2) for the duration of a 'with' statement.
If shared is True, use a LOCK_SH lock, otherwise LOCK_EX.
The file is given by name and will be created if it does not exist."""
with open(filename, "w") as lock, flock(lock, shared=shared):
yield
@contextmanager
def lockfile_nonblocking(filename: str) -> Iterator[bool]: # nocoverage
"""Lock a file using flock(2) for the duration of a 'with' statement.
Doesn't block, yields False immediately if the lock can't be acquired."""
with open(filename, "w") as f:
lock_acquired = False
try:
fcntl.flock(f, fcntl.LOCK_EX | fcntl.LOCK_NB)
lock_acquired = True
yield lock_acquired
except BlockingIOError:
yield False
finally:
if lock_acquired:
fcntl.flock(f, fcntl.LOCK_UN)