How to type a decorator that's used for methods of multiple subclasses in Python

Question

The code is like such, I tried to type g, which is to be used as a decorator for methods f of subclasses of A. f has predefined interface, so I want to avoid using typing.Any

from typing import Callable, Type, Union


def g(func: Callable[[A], int]) -> Callable[[A], int]:
    # do stuff and return a new callable.
    ...

class A:
    def __init__(self) -> None:
        self.a = 1

class B(A):
    def __init__(self) -> None:
        self.a = 2

    @g
    def f(self) -> int:
        return self.a


class C(A):
    def __init__(self) -> None:
        self.a = 3

    @g
    def f(self) -> int:
        return self.a

The above gave this error

mypy tmp.py
tmp.py:14:6: error: Argument 1 to "g" has incompatible type "Callable[[B], int]"; expected "Callable[[A], int]"
tmp.py:23:6: error: Argument 1 to "g" has incompatible type "Callable[[C], int]"; expected "Callable[[A], int]"
Found 2 errors in 1 file (checked 1 source file)

What's the proper way to type g in such case?

Answer 1

TypeVar can be used for your case:

from typing import Callable, TypeVar

T = TypeVar('T', bound='A')

def g(func: Callable[[T], int]) -> Callable[[T], int]:
    # do stuff and return a new callable.
    ...

class A:
    def __init__(self) -> None:
        self.a = 1

class B(A):
    def __init__(self) -> None:
        self.a = 2

    @g
    def f(self) -> int:
        return self.a


class C(A):
    def __init__(self) -> None:
        self.a = 3

    @g
    def f(self) -> int:
        return self.a