Path lock using a read/write lock on each ancestor path

path_lock_local_read_write.py

import asyncio
import contextlib
import weakref

from fifolock import FifoLock

class Read(asyncio.Future):
    @staticmethod
    def is_compatible(holds):
        return not holds[Write]

class Write(asyncio.Future):
    @staticmethod
    def is_compatible(holds):
        return not holds[Read] and not holds[Write]

class PathLock():
    def __init__(self):
        self._locks = weakref.WeakValueDictionary()

    @staticmethod
    def _sort_key(path_lock):
        return len(path_lock[0].parents), path_lock[0].as_posix()

    @staticmethod
    def _flatten(to_flatten):
        return [
            item for sub_list in to_flatten for item in sub_list
        ]

    def _with_locks(self, paths, mode):
        return [
            (path, self._locks.setdefault(path, default=FifoLock()), mode)
            for path in paths
        ]

    @contextlib.asynccontextmanager
    async def __call__(self, read, write):
        write_paths = set(read).union(set(write))
        write_locks = self._with_locks(write_paths, Write)

        read_paths = set(self._flatten(path.parents for path in write_paths)) - write_paths
        read_locks = self._with_locks(read_paths, Read)

        sorted_locks = sorted(read_locks + write_locks, key=self._sort_key)
        async with contextlib.AsyncExitStack() as stack:
            for _, lock, mode in sorted_locks:
                await stack.enter_async_context(lock(mode))

            yield