synchronised.py

from asyncio import Lock
from typing import Callable, Coroutine, ParamSpec, TypeVar, cast, overload


_locks: dict[str, Lock] = {}

P = ParamSpec("P")
T = TypeVar("T")


@overload
def synchronised(name: str) -> Callable[[Callable[P, Coroutine[None, None, T]]], Callable[P, Coroutine[None, None, T]]]:
    ...

@overload
def synchronised(name: Callable[P, Coroutine[None, None, T]]) -> Callable[P, Coroutine[None, None, T]]:
    ...


def synchronised(
    name: str | Callable[P, Coroutine[None, None, T]]
) -> Callable[[Callable[P, Coroutine[None, None, T]]], Callable[P, Coroutine[None, None, T]]] | Callable[P, Coroutine[None, None, T]]:
    if isinstance(name, str):
        def inner(func: Callable[P, Coroutine[None, None, T]]) -> Callable[P, Coroutine[None, None, T]]:
            nonlocal name
            name = cast(str, name)
            _locks[name] = Lock()

            async def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
                nonlocal name
                name = cast(str, name)
                await _locks[name].acquire()
                value = await func(*args, **kwargs)
                _locks[name].release()
                return value
            return wrapper
        return inner

    func = name
    name = name.__name__

    name = name or func.__name__
    _locks[name] = Lock()

    async def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
        nonlocal name
        name = cast(str, name)
        await _locks[name].acquire()
        value = await func(*args, **kwargs)
        _locks[name].release()
        return value
    return wrapper