1from __future__ import absolute_import
2from __future__ import unicode_literals
3
4import contextlib
5import errno
6
7
8try:  # pragma: no cover (windows)
9    import msvcrt
10
11    # https://docs.microsoft.com/en-us/cpp/c-runtime-library/reference/locking
12
13    # on windows we lock "regions" of files, we don't care about the actual
14    # byte region so we'll just pick *some* number here.
15    _region = 0xffff
16
17    @contextlib.contextmanager
18    def _locked(fileno, blocked_cb):
19        try:
20            msvcrt.locking(fileno, msvcrt.LK_NBLCK, _region)
21        except IOError:
22            blocked_cb()
23            while True:
24                try:
25                    msvcrt.locking(fileno, msvcrt.LK_LOCK, _region)
26                except IOError as e:
27                    # Locking violation. Returned when the _LK_LOCK or _LK_RLCK
28                    # flag is specified and the file cannot be locked after 10
29                    # attempts.
30                    if e.errno != errno.EDEADLOCK:
31                        raise
32                else:
33                    break
34
35        try:
36            yield
37        finally:
38            # From cursory testing, it seems to get unlocked when the file is
39            # closed so this may not be necessary.
40            # The documentation however states:
41            # "Regions should be locked only briefly and should be unlocked
42            # before closing a file or exiting the program."
43            msvcrt.locking(fileno, msvcrt.LK_UNLCK, _region)
44except ImportError:  # pragma: windows no cover
45    import fcntl
46
47    @contextlib.contextmanager
48    def _locked(fileno, blocked_cb):
49        try:
50            fcntl.flock(fileno, fcntl.LOCK_EX | fcntl.LOCK_NB)
51        except IOError:  # pragma: no cover (tests are single-threaded)
52            blocked_cb()
53            fcntl.flock(fileno, fcntl.LOCK_EX)
54        try:
55            yield
56        finally:
57            fcntl.flock(fileno, fcntl.LOCK_UN)
58
59
60@contextlib.contextmanager
61def lock(path, blocked_cb):
62    with open(path, 'a+') as f:
63        with _locked(f.fileno(), blocked_cb):
64            yield
65