Search code examples
pythontype-hintingpython-typingmagic-methodspython-descriptors

How to type hint python magic __get__ method


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)


Solution

  • 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]