How can I annotate a method that returns self so that the stub file shows correct return types?

user_7402287 0
Viewed 180

Pretext: This code is useless, and the use case might be useless also, but I am just trying to figure out how to make it work so that the stub files are showing the correct signatures. The question seems confusing, so I am trying my best to explain, but am very open to suggestions to make the question structure better.

Environment: Python: 3.7.9 Mypy Vscode + Pylance

Directory structure:

 .
├──  1.py
└──  testcode
   ├──  __init__.py
   └──  modules
      ├──  __init__.py
      └──  something.py

Inside 1.py

from testcode import Entry

e = Entry('a').a_function().b_function().c_function()

Inside __init__.py

from .modules.something import Something, AnotherThing

class Entry(Something, AnotherThing):
    def __init__(self, a: str):
        self.a = a

In this case, the return self is the instance of Entry. When I do not specify the return type in b_function (or a, b, c function), the type hints are correct.

enter image description here

But when I specify the return type as Entry, the return type shows as Any

something.py with return types annotated:

class Something(object):
    def a_function(self) -> 'Entry':
        return self

    def b_function(self) -> 'Entry':
        return self

class AnotherThing:
    def c_function(self) -> 'Entry':
        return self

enter image description here

So to try and overcome this, I tried creating stub files using stubgen with the command stubgen -o . -p testcode which creates all the pyi files, and something.pyi includes:

class Something:
    def a_function(self) -> Entry: ...
    def b_function(self) -> Entry: ...

class AnotherThing:
    def c_function(self) -> Entry: ...

Directory after stub generation:

 .
├──  1.py
└──  testcode
   ├──  __init__.py
   ├──  __init__.pyi
   └──  modules
      ├──  __init__.py
      ├──  __init__.pyi
      ├──  something.py
      └──  something.pyi

Now I understand why this doesnt work, because when i specify the return type of Entry in a method for Something, python doesnt know what Entry is which can be seen in the next screenshot. So it assigns the Any type to it:

enter image description here

Question:

Knowing that when the return type of Entry is not annotated in a method for Something, it picks up the correct return type (and hence provides nice intellisense), how can I annotate and generate stubfiles correctly so that it doesnt return Any when a return type of self is annotated?

Edit:

To reflect on the first answer, binding the subclasses does not fully solve the problem because the intellisense and types will show only the called methods methods. So for example, if I call a_function, c_function will not show, and vice versa. When I dont type hint on the return types of the Something and AnotherThing class, and vscode correctly assumes Entry is the return type, then i can get intellisense for c_function when I call a_function

enter image description here

python-typing mypy python
1 Answers

The solution is to make the self parameter of your Something and AnotherThing functions bound to a subclass of those classes using a TypeVar like so:

from typing import TypeVar


ST = TypeVar('ST', bound='Something')

class Something:
    def a_function(self: ST) -> ST:
        return self

    def b_function(self: ST) -> ST:
        return self

AT = TypeVar('AT', bound='AnotherThing')

class AnotherThing:
    def c_function(self: AT) -> AT:
        return self

More info can be found here.

user_113542660
Related