How to type hint python magic __get__ method

420 Views Asked by At

Suppose we have the following classes:

class Foo:
   def __init__(self, method):
       self.method = method

   def __get__(self, instance, owner):
       if instance is None:
          return self
       return self.method(instance)


class Bar:
    @Foo
    def do_something(self) -> int:
        return 1


Bar().do_something  # is 1
Bar.do_something    # is Foo object

How to type hint __get__ and method correctly so that Pylance understands Bar().do_something is of the return type of do_something? (like a standard property)

1

There are 1 best solutions below

5
Daniil Fajnberg On BEST ANSWER

You'll need to overload the __get__ method.

I do not use VSCode myself, but I tested the code below with MyPy and I would expect Pyright to infer the types correctly as well.

Python >=3.9

To make this as flexible as possible, I would suggest making Foo generic in terms of

  1. the class using the descriptor/decorator,
  2. the parameter specification of the decorated method, and
  3. the return type of the decorated method.
from collections.abc import Callable
from typing import Generic, TypeVar, Union, overload
from typing_extensions import Concatenate, ParamSpec, Self

T = TypeVar("T")    # class using the descriptor
P = ParamSpec("P")  # parameter specs of the decorated method
R = TypeVar("R")    # return value of the decorated method


class Foo(Generic[T, P, R]):
    method: Callable[Concatenate[T, P], R]

    def __init__(self, method: Callable[Concatenate[T, P], R]) -> None:
        self.method = method

    @overload
    def __get__(self, instance: T, owner: object) -> R: ...

    @overload
    def __get__(self, instance: None, owner: object) -> Self: ...

    def __get__(self, instance: Union[T, None], owner: object) -> Union[Self, R]:
        if instance is None:
            return self
        return self.method(instance)

Demo:

from typing import TYPE_CHECKING


class Bar:
    @Foo
    def do_something(self) -> int:
        return 1


a = Bar().do_something
b = Bar.do_something

print(type(a), type(b))  # <class 'int'> <class '__main__.Foo'>
if TYPE_CHECKING:
    reveal_locals()

Running MyPy over this gives the desired output:

note: Revealed local types are:
note:     a: builtins.int
note:     b: Foo[Bar, [], builtins.int]

NOTE: (thanks to @SUTerliakov for pointing some of this out)

  • If you are on Python >=3.10, you can import Concatenate and ParamSpec directly from typing and you can use the |-notation instead of typing.Union.
  • If you are on Python >=3.11 you can import Self directly from typing as well, meaning you won't need typing_extensions at all.

Python <3.9

Without Concatenate, ParamSpec and Self we can still make Foo generic in terms of the return value of the decorated method:

from __future__ import annotations
from collections.abc import Callable
from typing import Generic, TypeVar, Union, overload

R = TypeVar("R")    # return value of the decorated method


class Foo(Generic[R]):
    method: Callable[..., R]

    def __init__(self, method: Callable[..., R]) -> None:
        self.method = method

    @overload
    def __get__(self, instance: None, owner: object) -> Foo[R]: ...

    @overload
    def __get__(self, instance: object, owner: object) -> R: ...

    def __get__(self, instance: object, owner: object) -> Union[Foo[R], R]:
        if instance is None:
            return self
        return self.method(instance)

MyPy output for the same demo script from above:

note: Revealed local types are:
note:     a: builtins.int
note:     b: Foo[builtins.int]