Type hinting a decorator that adds a parameter with default value

Question

Motivation: I realised I had a lot of class methods that were also being used as TKinter callbacks, which pass a tk.Event object as the first (non-self) argument. If my application wants to call them normally as well, this event argument should be None by default...

class Writer:
    # very unwieldy second parameter...
    def write(self, event: Optional[tk.Event] = None, number: int = 0) -> str:
        return str(number)

It's more boilerplate, it forces me to provide default arguments for everything, and pylint is screaming about unused arguments.

So I wrote a decorator to add the extra parameter... but how do I type hint it correctly? (I'm using Python 3.8.10 and mypy 0.971.)

def tk_callback(method):
    @functools.wraps(method)
    def wrapper(self, event=None, *args, **kwargs):
        return method(self, *args, **kwargs)
    return wrapper

The callback should not hide the types of the original parameters. It should reflect that the added parameter event gets passed a default value of None.

I did some reading about generics, Protocols, and a little searching (e.g.)

• How to decorate a function in python in a way that adds a new argument and creates the correct signature?
• Python decorator adding argument to function and its signature

The linked questions are similar but not duplicates: I'd like type hinting, and specifically make the extra argument on the wrapper function take a default value. Then I made an attempt:

# foo.py
from __future__ import annotations

import functools
import tkinter as tk

from typing import Callable, Optional, Protocol, TypeVar
from typing_extensions import Concatenate, ParamSpec


P = ParamSpec("P")
CT_contra = TypeVar("CT_contra", contravariant=True)
RT_co = TypeVar("RT_co", covariant=True)
C = TypeVar("C")
R = TypeVar("R")


class Prot(Protocol[CT_contra, P, RT_co]):
    def __call__(self,
            _: CT_contra, # this would be "self" on the method itself
            event: Optional[tk.Event] = ..., /,
            *args: P.args,
            **kwargs: P.kwargs
        ) -> RT_co:
        ...


def tk_callback(method: Callable[Concatenate[C, P], R]
    ) -> Prot[C, P, R]:
    @functools.wraps(method)
    def wrapper(
            self: C,
            event: Optional[tk.Event] = None,
            *args: P.args,
            **kwargs: P.kwargs
        ) -> R:
        return method(self, *args, **kwargs)
    return wrapper

Which doesn't seem to work. mypy complains the decorator doesn't return what the type hint declares. error: Incompatible return value type (got "Callable[[C, Optional[Event[Any]], **P], R]", expected "Prot[C, P, R]")

It also notes that the returned wrapper functions should have a very similar type: "Prot[C, P, R].__call__" has type "Callable[[C, DefaultArg(Optional[Event[Any]]), **P], R]"

(Digression: not relevant to my use case, but if I don't supply the default argument in the protocol, it still complains while noting that "Prot[C, P, R].__call__" has type "Callable[[C, Optional[Event[Any]], **P], R]") even though this is exactly what is returned!)

So what should be the right way to type hint this decorator, and/or how can I get the type checking to work correctly?

---

More troubleshooting information: the revealed type of a method is also strange.

# also in foo.py
class Writer:
    def __init__(self) -> None:
        return

    @tk_callback
    def write(self, number: int) -> str:
        return str(number)

writer = Writer()
reveal_type(writer.write) # mypy: Revealed type is "foo.Prot[foo.Writer, [number: builtins.int], builtins.str]

Answer 1

This question boils down to:

How do I type hint a decorator which adds an argument to a class method?

Inspired by the docs for typing.Concatenate, I've come up with a solution which works when the first parameter of the function being decorated is self.

I've simplified your example a bit, to use an optional string instead of your more involved tk.Event which isn't really relevant to the question.

from collections.abc import Callable
from typing import Concatenate, Optional, ParamSpec, TypeVar


S = TypeVar("S")
P = ParamSpec("P")
R = TypeVar("R")


def with_event(
    method: Callable[Concatenate[S, P], R]
) -> Callable[Concatenate[S, Optional[str], P], R]:
    def wrapper(
        _self: S, event: Optional[str] = None, *args: P.args, **kwargs: P.kwargs
    ) -> R:
        print(f"event was {event}")
        return method(_self, *args, **kwargs)

    return wrapper


class Writer:
    @with_event
    def write(self, number: int = 0) -> str:
        return f"number was {number}"


writer = Writer()
print(writer.write("my event", number=5))

This prints:

event was my event
number was 5

The with_event signature works like this:

def with_event(
    # Take one parameter which is a function with takes a
    # single positional argument (S), then some other stuff (P)
    # and return another function which takes the same position
    # argument S, followed by an Optional[str], followed by
    # the "other stuff", P. The return types of these two functions
    # is the same (R).
    method: Callable[Concatenate[S, P], R]
) -> Callable[Concatenate[S, Optional[str], P], R]: