#+TITLE: 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:

#+BEGIN_SRC python
  from threading import RLock

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

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

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:

#+BEGIN_SRC python
  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: ...
#+END_SRC

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...

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

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.

* 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 [[https://mypy.readthedocs.io/en/stable/generics.html#declaring-decorators][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):

#+BEGIN_SRC python
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 []
#+END_SRC

This failed with:

#+BEGIN_EXAMPLE
  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
#+END_EXAMPLE

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

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

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=:

#+BEGIN_SRC python
  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]"
#+END_SRC

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:

#+BEGIN_SRC python
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
#+END_SRC