id_lock.py

import asyncio
from collections import defaultdict
from collections.abc import AsyncIterator, Hashable
from contextlib import asynccontextmanager
from typing import Self


class IDLock:
    """Asyncio Lock synced to some id."""

    def __init__(self: Self) -> None:
        self._lock_dict: dict[Hashable, asyncio.Lock] = defaultdict(asyncio.Lock)

    @asynccontextmanager
    async def ctx(self: Self, any_id: Hashable) -> AsyncIterator[None]:
        """Lock for any id.

        Usage:
            lock = IDLock()  # once per app, not per request!
            
            # once per request
            async with lock.ctx(user_id):
                ...
        """

        lock = self._lock_dict[any_id]
        try:
            async with lock:
                yield
        finally:
            if not lock._waiters and any_id in self._lock_dict:  # noqa till issue:
                # https://github.com/python/typeshed/issues/11314 
                # is not imported to pyCharm
                del self._lock_dict[any_id]