CODEWITHSUNDEEP
pythonpython-typing
wakey
wakey

Reputation: 2409

Correctly typehint functions returning self in python3

Say I have the following code....

from __future__ import annotations
from typing import Any, TypeVar, Type
from typing_extensions import Literal

_T = TypeVar("_T")

class A:

    @classmethod
    def constructor(cls: Type[_T]) -> _T:
        return cls()

    def __enter__(self) -> A:
        return self

    def __exit__(self, *args: Any, **kwargs: Any) -> Literal[False]:
        return False


class B(A):
    pass


def test() -> None:
    a_instance: A = A.constructor()
    b_instance: B = B.constructor()  # Works because constructor return type is _T

    with B() as b:
        a_instance = b
        b_instance = b  # Error, Incompatible types in assignment

If I run mypy against the above code, I get the following warning

» mypy test.py                               
test.py:30: error: Incompatible types in assignment (expression has type "A", variable has type "B")
Found 1 error in 1 file (checked 1 source file)

This is because the return type of A.__enter__, (A) is inherited by B, so mypy thinks that B.__enter__ also returns A. I'd like to have to avoid re-implementing the function just to correct the typehint...

class B(A):
   def __enter__(self) -> B:
      return super().__enter__()  # type: ignore

I've gotten around similar issues in classmethod constructors by using a TypeVar to template away the type of the cls so I can use it later, but I'm not sure how a similar trick could apply to non-classmethods.

This situation also applies for any other method that returns self.

Upvotes: 0

Views: 98

Answers (1)

wakey
wakey

Reputation: 2409

Believe I figured this out...

You need to also use the TypeVar _T to indicate the type of self. Then you can re-use that same _T as the return type.

Note, if you want to use any attributes of the class (like I do in the print below), you also need to bind the TypeVar to a class using the bound argument.

_T = TypeVar("_T", bound="A")

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

    @classmethod
    def constructor(cls: Type[_T]) -> _T:
        return cls()

    def __enter__(self: _T) -> _T:
        print(self.a_attr)
        return self

Upvotes: 1

Related Questions