Defining a type-checked `@synchronise` method decorator in Python

I have a class defining a long-lived object, and the methods on that object may be called from many threads, so those that modify the inner state, or those that are more complex than merely returning a field, need some locking.

That's quite easy, even if it turned into a bit of a rabbit-hole.

First version:

from threading import RLock

class Foo:
    def __init__(self) -> None:
        self.lock = RLock()

    def do_something(self, arg: int) -> str:
        with self.lock:
            ...

Simple, to the point. However, that pushes the bodies of all functions that need locking one indentation level to the right, and I found it's easy to forget the with block. We can make a decorator:

def synchronise(meth):
    def wrapper(self, *args, **kwargs):
        with self.lock:
            meth(self, *args, **kwargs)
    return wrapper

class Foo:
    @synchronise
    def do_something(self, arg: int) -> str: ...

It works just fine. We save one level of indentation, and I think adding decorators to methods that need it is less error-prone than having a with statement that spans the entire method body. If we need finer-grained locking, we can leave out the @synchronise decorator and lock only the critical section. However…

reveal_type(Foo().do_something) # main.py:11: note: Revealed type is "Any"

The type of Foo.do_something is completely erased! If you're not bothered about mypy-style static type hinting, maybe that's good enough for you, and you can stop reading. As for myself, I've found those to be immensely useful, so let's try to fix that.

1. Type-hinting a decorator

The reason for Foo.do_something becoming Any is that synchronise is missing any type annotations, so it's return type (which is installed as Foo.do_something) is Any. Thanksfull mypy documentation has a chapter about annotating decorators. The common case of a decorator that doesn't care (or change) the types of the callable arguments is simple enough; but we want to ensure the first arguments of our method (the self) actually has a .lock member of the right type.

My first attempt (using the legacy syntax (Python 3.11 and earlier), as I still have machines running Debian bookworm with Python 3.11):

from threading import RLock
from typing import Callable, Concatenate, ParamSpec, Protocol, TypeVar


class Lockable(Protocol):
    lock: RLock


P = ParamSpec('P')
# P_lockable = Concatenate[Lockable, P]
T = TypeVar('T')


def synchronise(meth: Callable[Concatenate[Lockable, P], T]) -> Callable[Concatenate[Lockable, P], T]:
    def wrapper(self: Lockable, *args: P.args, **kwargs: P.kwargs) -> T:
        with self.lock:
            return meth(self, *args, **kwargs)
    return wrapper


class Foo(Lockable):
    def __init__(self) -> None:
        self.lock = RLock()

    @synchronise
    def do_something(self) -> list[str]:
        return []

This failed with:

main.py:24: error: Argument 1 to "synchronise" has incompatible type "Callable[[Foo], list[str]]"; expected "Callable[[Lockable], list[str]]"  [arg-type]
main.py:24: note: This is likely because "do_something of Foo" has named arguments: "self". Consider marking them positional-only

Changing the annotations of the decorator to be exactly Foo make mypy happy:

def synchronise(meth: Callable[Concatenate[Foo, P], T]) -> Callable[Concatenate[Foo, P], T]:
    def wrapper(self: Foo, *args: P.args, **kwargs: P.kwargs) -> T:
        ...

That tells us the "note" about self because a named argument is a red herring, and gives us a hint, that the problem is the type of the first argument. The issue is that Callable is contravariant with respect to arguments. In the previous version, if meth can accept any Lockable as first argument, then Foo.do_something which accepts only Foo is less generic than meth, so can't be accepted.

That's okay, we can rewrite this with a type variable bound to Lockable:

from threading import RLock
from typing import Callable, Concatenate, ParamSpec, Protocol, TypeVar

class Lockable(Protocol):
    lock: RLock

L = TypeVar('L', bound=Lockable)
P = ParamSpec('P')
T = TypeVar('T')

def synchronise(meth: Callable[Concatenate[L, P], T]) -> Callable[Concatenate[L, P], T]:
    def wrapper(self: L, *args: P.args, **kwargs: P.kwargs) -> T:
        with self.lock:
            return meth(self, *args, **kwargs)
    return wrapper

class Foo:
    def __init__(self) -> None:
        self.lock = RLock()

    @synchronise
    def do_something(self) -> list[str]:
        return []

reveal_type(Foo().do_something) # note: Revealed type is "def () -> builtins.list[builtins.str]"

It all works! The method keeps the proper type; so any misuse will raise a type error.. If you try to apply @synchronise to a method in a class that doesn't define a .lock member (or of the wrong type), you get a type error. A little bit heavy on the type programming, but very simple usage.

Bonus, the syntax for Python 3.12+ allows us to skip defining L, P, T:

from threading import Lock, RLock
from typing import Callable, Concatenate, Protocol

class Lockable(Protocol):
    lock: RLock

def synchronise[L: Lockable, **P, T](meth: Callable[Concatenate[L, P], T]) -> Callable[Concatenate[L, P], T]:
    def wrapper(self: L, *args: P.args, **kwargs: P.kwargs) -> T:
        with self.lock:
            return meth(self, *args, **kwargs)
    return wrapper