Possible to mypy typehint a method that takes a `cls` argument rather than `self`?

Viewed 83

I'm curious if it's possible to get mypy to understand that some method in a class takes a cls argument rather than a self argument. The concrete use-case I have in mind is typing a sqlalchemy hybrid_property's .expression decorator, but a simpler example would be trying to implement a custom/non-built-in version of classmethod.

Applying an explicit type annotation to cls doesn't work, e.g.

from typing import Type

class Example:
  @my_custom_classmethod
  def foo(cls: Type["Example"]) -> str:
    return "doesn't work"

Mypy complains that The erased type of self "Type[Example]" is not a supertype of its class "Example", which, ok, I guess often makes sense, but it doesn't make sense here :/

The best solution I've managed to come up with so far is to just # type: ignore the method's signature:

class Example:
  @my_custom_classmethod
  def foo(cls: Type["Example"]) -> str: # type: ignore
    return "this actually kind of does work!"

This actually does seem to basically work, insofar as cls has the right type in the body of the method and the decorator receives a method of the right type, but it feels a bit hacky.

1 Answers

bind_self is a method that checks and assigns self (cls) type. It relies on is_classmethod node flag (implementation here).

Flag is_classmethod is set during semantic analysis here. It is done by direct comparison of decorator and builtins.classmethod. So your custom_classmethod is not considered classmethod.

So you cannot produce decorator that will be equivalent to classmethod builtin. However, you can overcome this problem using custom descriptor:

from __future__ import annotations
from typing import Callable, Concatenate, Generic, ParamSpec, TypeVar, TYPE_CHECKING

_T = TypeVar('_T')
_R = TypeVar('_R')
_P = ParamSpec('_P')


if TYPE_CHECKING:
    base = classmethod
else:
    class base(Generic[_R], classmethod):
        pass
    
class my_classmethod(base[_R], Generic[_P, _R]):
    @property
    def __func__(self) -> Callable[Concatenate[_T, _P], _R]:
        return super().__func__
    def __init__(self: my_classmethod[_P, _R], f: Callable[Concatenate[_T, _P], _R], /) -> None:
        super().__init__(f)
    def __get__(self, obj: _T, type: type[_T] | None = None) -> Callable[_P, _R]:
        return super().__get__(obj, type)
        
    
class A:
    @my_classmethod
    def foo(cls, x: int) -> str:
        return str(x)

reveal_type(A.foo(1))  # N: Revealed type is "builtins.str"
reveal_type(A().foo(1))  # N: Revealed type is "builtins.str"
A.foo('a')  # E: Argument 1 has incompatible type "str"; expected "int"

This works on runtime and passes type checking. Descriptors are funny. I'm assuming python >= 3.10 and recent enough mypy: older python versions may require some changes (at least imports from typing_extensions), old mypy can have troubles with full descriptor support. And yes, you can inherit from classmethod. If you don't want, you'll have to re-implement __get__ and __init__ anyway (and __func__ property if you like consistency). Remember that in 3.10 and later classmethod also has wrapped property.

Here's playground with this code.

Related