I am trying to make a system that records the methods calls of a class. (I already encountered some issues for which I posted yesterday and they were solved).

The idea is the following:

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

RetType = TypeVar("RetType")
Param = ParamSpec("Param")
CountType = TypeVar("CountType", bound="FunctionCount")


class FunctionCount:
    def __init__(self, count_dict: dict[str, int]) -> None:
        self.count_dict = count_dict


def count(
    func: Callable[Concatenate[CountType, Param], RetType],
) -> Callable[Concatenate[CountType, str, Param], RetType]:
    def wrapper(
        self: CountType,
        /,
        *args: Param.args,
        log_name: str | None = None,
        **kwargs: Param.kwargs,
    ) -> RetType:
        function_name = log_name or f"{self.__class__.__name__}.{func.__name__}"
        if function_name not in self.count_dict:
            self.count_dict[function_name] = 0
        self.count_dict[function_name] += 1

        return func(self, *args, **kwargs)

    return wrapper


class A(FunctionCount):
    @count
    def test(self, multiplier: float) -> float:
        return multiplier


count_dict = dict[str, int]()
A(count_dict).test(2.0)
A(count_dict).test(multiplier=2.0, log_name="testing")

print(count_dict)
assert count_dict == {"A.test": 1, "testing": 1}

The script is executed without any problem but mypy has problems
It signals me that wrapper does not have the correct type:

Incompatible return value type (got "Callable[[CountType, VarArg(Any), DefaultNamedArg(str | None, 'log_name'), KwArg(Any)], RetType]", expected "Callable[[CountType, str, **Param], RetType]")

I guess this is because I wanted log_name to be a keyword-only argument so understanding the call A(count_dict).test(2.0, log_name="testing") is quite easier than if we had A(count_dict).test(2.0, "testing")

I tried to use Protocol for the typing but then I had other problems

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

RetType_co = TypeVar("RetType_co", covariant=True)
Param = ParamSpec("Param")
CountType_contra = TypeVar(
    "CountType_contra", bound="FunctionCount", contravariant=True
)


class FunctionCount:
    def __init__(self, count_dict: dict[str, int]) -> None:
        self.count_dict = count_dict


class CountableCallable(Protocol[CountType_contra, Param, RetType_co]):
    def __call__(
        _self,  # noqa: N805
        self: CountType_contra,
        /,
        *args: Param.args,
        log_name: str | None = None,
        **kwargs: Param.kwargs,
    ) -> RetType_co: ...


def count(
    func: Callable[Concatenate[CountType_contra, Param], RetType_co],
) -> CountableCallable[CountType_contra, Param, RetType_co]:
    def wrapper(
        self: CountType_contra,
        /,
        *args: Param.args,
        log_name: str | None = None,
        **kwargs: Param.kwargs,
    ) -> RetType_co:
        function_name = log_name or f"{self.__class__.__name__}.{func.__name__}"
        if function_name not in self.count_dict:
            self.count_dict[function_name] = 0
        self.count_dict[function_name] += 1

        return func(self, *args, **kwargs)

    return wrapper


class A(FunctionCount):
    @count
    def test(self, multiplier: float) -> float:
        return multiplier


count_dict = dict[str, int]()
A(count_dict).test(2.0)
A(count_dict).test(multiplier=2.0, log_name="testing")

print(count_dict)
assert count_dict == {"A.test": 1, "testing": 1}

In this case, the A(count_dict).test(2.0) lines raise the following error:

Argument 1 to "__call__" of "CountableCallable" has incompatible type "float"; expected "A"

Calling A(count_dict).test(multiplier=2.0, log_name="testing") raises this error:

Missing positional argument "self" in call to "__call__" of "CountableCallable"

The only calling working is A.test(A(count_dict), 2.0) which is not the way people call methods.

I also tried to use a decorator instead of a wrapper but I would get the same alerts
I tried to intervert the names of the self and __self__ arguments for CountableCallable.__call__ but it did not do anything

So I was wondering if there was a Protocol variation that would allow me to specify that CountableCallable is a class method (so maybe we wouldn’t have confusion)
Or if there was a way to use typing.Callable while precising that log_name is a keyword-only argument