Mypy: how to use type aliases in generic classes

[zultron]

I'm working on a minimal example (in Py3.8) of a decorator with these qualities:

• Implemented in a class
• Decorates class methods
• Can be used with or without args
• Passes mypy --strict

I finally came up with this (I'm new to Python type hinting, needed a little help from STerliakov on SO to get me across the finish line):

#!/usr/bin/env python3.8
from __future__ import annotations
from typing import (
    Callable,
    TypeVar,
    Generic,
    Any,
    cast,
    overload,
)
from typing_extensions import Concatenate, ParamSpec, reveal_type

# TypeVars for decorator without args:
T_o_wo = TypeVar("T_o_wo")  # Type of the class instance
P_wo = ParamSpec("P_wo")  # Param spec for decorated unbound function arguments
T_ret_wo = TypeVar("T_ret_wo")  # Return type of the decorated function

# TypeVars for decorator with args:
T_o_w = TypeVar("T_o_w")  # Type of the class instance
P_w = ParamSpec("P_w")  # Param spec for decorated unbound function arguments
T_ret_w = TypeVar("T_ret_w")  # Return type of the decorated function

# *Type aliases in generics break mypy
# T_c_wo = Callable[Concatenate[T_o_wo, P_wo], T_ret_wo]  # Type of funct
# T_c_w = Callable[Concatenate[T_o_w, P_w], T_ret_w]  # Type of funct

class MyDecoCls(Generic[T_o_wo, T_ret_wo, P_wo]):

    @overload  # Decorator with args:  __init__(self, arg1=..., arg2=...)
    def __init__(self, /, arg1: bool = False, arg2: bool = False) -> None: ...

    @overload  # Decorator without args:  __init__(self, funct)
    # *Type aliases in generics break mypy
    # def __init__(self, funct: T_c_wo) -> None: ...
    def __init__(
        self, funct: Callable[Concatenate[T_o_wo, P_wo], T_ret_wo]
    ) -> None: ...

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        self.with_args = not (args and callable(args[0]))
        print(f"In __init__({self}, *{args}, **{kwargs})")
        # *Type aliases in generics break mypy
        # self.funct: T_c_wo | None = None
        self.funct: Callable[Concatenate[T_o_wo, P_wo], T_ret_wo] | None = None
        self.arg1 = kwargs.get("arg1", False)
        self.arg2 = kwargs.get("arg2", False)
        if not self.with_args:
            # *Type aliases in generics break mypy
            # self.funct = cast(T_c_wo, args[0])
            self.funct = cast(
                Callable[Concatenate[T_o_wo, P_wo], T_ret_wo], args[0]
            )

    def __get__(
        self, instance: T_o_wo, owner: type[T_o_wo]
    ) -> Callable[P_wo, T_ret_wo]:
        # Decorator without args:  __get__(self, obj, type(obj)) -> wrapper
        print(f"In __get__({self}, {instance}, {owner})")
        assert instance
        assert self.funct

        def wrapper(*args: P_wo.args, **kwargs: P_wo.kwargs) -> T_ret_wo:
            assert self.funct
            ret = self.funct(instance, *args, **kwargs)  # Use instance here
            print(f"In wrapper(*{args}, **{kwargs}) -> {ret} (__get__)")
            print(f"    self.funct: {self.funct}; instance: {instance}")
            return ret

        reveal_type(wrapper)
        return wrapper

    reveal_type(__get__)

    # *Type aliases in generics break mypy
    # def __call__(self, funct: T_c_w) -> T_c_w:
    def __call__(
        self, funct: Callable[Concatenate[T_o_w, P_w], T_ret_w]
    ) -> Callable[Concatenate[T_o_w, P_w], T_ret_w]:
        # Decorator with args:  __call__(self, funct) -> wrapper
        print(f"In __call__({self}, {funct})")
        assert self.with_args
        assert callable(funct)
        # self.funct = funct  # type: ignore[assignment]

        def wrapper(
            self: T_o_w, /, *args: P_w.args, **kwargs: P_w.kwargs
        ) -> T_ret_w:
            ret = funct(self, *args, **kwargs)
            print(
                f"In wrapper({self}, *{args}, **{kwargs}) -> {ret}  (__call__)"
            )
            return ret

        reveal_type(wrapper)
        return wrapper

    reveal_type(__call__)

    def __str__(self) -> str:
        args = "with args" if self.with_args else "without args"
        return f"<MyDecoCls object ({args}) {hex(id(self))[-6:]}>"


# Example usage
class Foo:
    @MyDecoCls
    def bar(self, i: int) -> int:
        print(f"{self}.bar({i})")
        return i * 2

    reveal_type(bar)

    @MyDecoCls(arg1=True)
    def baz(self, /, i: int) -> int:
        print(f"{self}.baz(i={i})")
        return i * 4

    reveal_type(baz)

    def __str__(self) -> str:
        return "<Foo object>"


# Test the implementation using the example usage
f = Foo()
res: int
print(f"object f = {f}")
print(f"running bar: {(res := f.bar(42))}")
assert res == 42 * 2, f"res was {res}"
print(f"running baz: {(res := f.baz(i=13))}")
assert res == 13 * 4, f"res was {res} != 13 * 4"

This works fine, exactly what I need. It would be a little more readable if the repeeated, wordy Callable[] expressions could be replaced with aliases; note the locations marked with # *Type aliases in generics break mypy in the listing.

The above uses the generic class internals support in mypy, of course. With the commented-out type aliases, I'm trying to use the generic type aliases support, but get the following errors. What am I doing wrong?

Program run:

$ python3 foo.py
Runtime type is 'function'
Runtime type is 'function'
In __init__(<MyDecoCls object (without args) dc40a0>, *(<function Foo.bar at 0x7fd6d4c33e50>,), **{})
Runtime type is 'MyDecoCls'
In __init__(<MyDecoCls object (with args) d09a30>, *(), **{'arg1': True})
In __call__(<MyDecoCls object (with args) d09a30>, <function Foo.baz at 0x7fd6d4c33ee0>)
Runtime type is 'function'
Runtime type is 'function'
object f = <Foo object>
In __get__(<MyDecoCls object (without args) dc40a0>, <Foo object>, <class '__main__.Foo'>)
Runtime type is 'function'
<Foo object>.bar(42)
In wrapper(*(42,), **{}) -> 84 (__get__)
    self.funct: <function Foo.bar at 0x7fd6d4c33e50>; instance: <Foo object>
running bar: 84
<Foo object>.baz(i=13)
In wrapper(<Foo object>, *(), **{'i': 13}) -> 52  (__call__)
running baz: 52

mypy --strict run:

$ mypy --strict foo.py
foo.py:34: error: Missing type parameters for generic type "T_c_wo"  [type-arg]
foo.py:43: error: Missing type parameters for generic type "T_c_wo"  [type-arg]
foo.py:49: error: Missing type parameters for generic type "T_c_wo"  [type-arg]
foo.py:67: error: Returning Any from function declared to return "T_ret_wo"  [no-any-return]
foo.py:69: note: Revealed type is "def (*P_wo.args, **P_wo.kwargs) -> T_ret_wo`2"
foo.py:72: note: Revealed type is "def (self: foo.MyDecoCls[T_o_wo`1, T_ret_wo`2, P_wo`3], instance: T_o_wo`1, owner: Type[T_o_wo`1]) -> def (*P_wo.args, **P_wo.kwargs) -> T_ret_wo`2"
foo.py:75: error: Missing type parameters for generic type "T_c_w"  [type-arg]
foo.py:87: error: A function returning TypeVar should receive at least one argument containing the same TypeVar  [type-var]
foo.py:92: error: Returning Any from function declared to return "T_ret_w"  [no-any-return]
foo.py:94: note: Revealed type is "def [T_o_w, P_w, T_ret_w] (T_o_w`-1, *P_w.args, **P_w.kwargs) -> T_ret_w`-3"
foo.py:97: note: Revealed type is "def (self: foo.MyDecoCls[T_o_wo`1, T_ret_wo`2, P_wo`3], funct: def (Any, *Any, **Any) -> Any) -> def (Any, *Any, **Any) -> Any"
foo.py:106: error: Argument 1 to "MyDecoCls" has incompatible type "Callable[[Foo, int], int]"; expected "Callable[[Any, VarArg(Any), KwArg(Any)], Any]"  [arg-type]
foo.py:106: note: This is likely because "bar of Foo" has named arguments: "self". Consider marking them positional-only
foo.py:111: note: Revealed type is "foo.MyDecoCls[Never, Never, Never]"
foo.py:118: note: Revealed type is "def (Any, *Any, **Any) -> Any"
foo.py:128: error: Argument 1 to "__get__" of "MyDecoCls" has incompatible type "Foo"; expected "Never"  [arg-type]
foo.py:128: error: Argument 2 to "__get__" of "MyDecoCls" has incompatible type "Type[Foo]"; expected "Type[Never]"  [arg-type]
foo.py:128: error: Argument 1 has incompatible type "int"; expected "Never"  [arg-type]
Found 11 errors in 1 file (checked 1 source file)

[erictraut]

When you refer to a generic type alias in your code, you need to provide type arguments for it. If you don't provide type arguments, type checkers must (according to the typing spec) assume that you intended the type arguments to be Any.

In your code sample, you've defined the type alias T_c_wo as a generic type alias that has three generic type parameters. When you use T_c_wo, you should provide three type arguments. These type arguments can include type variables that are scoped to your local function or class. For example: x: T_c_wo[int, ..., str] or x: T_c_wo[T_o_wo, P_wo, T_ret_wo]. If you use T_c_wo by itself with non subscript, it will be interpreted as T_c_wo[Any, ..., Any]. That's probably not what you intend.

When a TypeVar is allocated as a global variable, it has no meaning until it is bound to a scope. A scope can be a generic class, a generic function, or a generic type alias. In your code sample, you're binding some of these type variable to multiple different scopes. For example, T_o_wo is bound to the scope of the type alias T_c_wo and also the class MyDecoCls. At runtime, T_o_wo is the same global object, but a type checker will treat these as completely independent type variables because they have different scopes.

If you're using Python 3.12, you may want to switch to the new PEP 695 syntax for defining generics. It eliminates much of this confusion and will feel much more natural if you're used to programming with generics in other languages. Mypy recently added experimental support for this new syntax.

[erictraut]

Here's a simplified version of your code using modern syntax and with all cast calls eliminated.

Code sample in pyright playground

from typing import Any, Callable, overload, Concatenate

type TCall[T, **P, R] = Callable[Concatenate[T, P], R]


class MyDecoCls[T, R, **P]:
    @overload
    def __init__(self, /, arg1: bool = False, arg2: bool = False) -> None: ...

    @overload
    def __init__(self, funct: TCall[T, P, R], /) -> None: ...

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        self.with_args = not (args and callable(args[0]))
        self.funct = None
        self.arg1 = kwargs.get("arg1", False)
        self.arg2 = kwargs.get("arg2", False)
        if not self.with_args:
            self.funct = args[0]

    def __get__(self, instance: T, owner: type[T]) -> Callable[P, R]:
        def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
            assert self.funct is not None
            return self.funct(instance, *args, **kwargs)

        return wrapper

    def __call__(self, funct: TCall[T, P, R]) -> TCall[T, P, R]:
        self.funct = funct

        def wrapper(self: T, /, *args: P.args, **kwargs: P.kwargs) -> R:
            return funct(self, *args, **kwargs)

        return wrapper

    def __str__(self) -> str:
        args = "with args" if self.with_args else "without args"
        return f"<MyDecoCls object ({args}) {hex(id(self))[-6:]}>"