RLock for asyncio

asyncio_rlock.py

from typing import Any
import asyncio


class RLock(asyncio.Lock):
    """ RLock like threading.RLock but use asyncio tasks as an ident """
    def __init__(self, *args: Any, **kwargs: Any):
        super().__init__(*args, **kwargs)
        self.loop = asyncio.get_event_loop()
        self._owner: Optional[asyncio.Task] = None
        self._count: int = 0

    def __get_ident(self) -> int:
        return id(asyncio.current_task(self.loop))

    async def acquire(self) -> bool:
        me = self.__get_ident()
        if self._owner == me:
            self._count += 1
            return True

        await super().acquire()
        self._owner = me
        self._count = 1
        return True

    def release(self) -> None:
        me = self.__get_ident()
        if self._owner != me:
            raise RuntimeError("cannot release un-acquired lock")

        self._count = count = self._count - 1
        if not count:
            self._owner = None
            super().release()
            

if __name__ == "__main__":
    async def test():
        lock = RLock()
        x = 1
        async with lock:
            x += 1
            async with lock:
              x += 1
        assert x == 3

    asyncio.run(test())