Can mypy select a method return type based on the type of the current object?

问题

In the following code, calling the clone() method on an instance of A will return an instance of type A, calling it on an instance of B will return an instance of type B, and so on. The purpose is to create a new instance which is identical to the current one but has a different internally generated primary key, so it can be edited from there and safely saved as a new item.

??? is some kind of type qualifier, but I'm not sure what the right choice is.

class A:
    def clone(self) -> ???:
        cls = self.__class__
        result = cls.__new__(cls)
        for k, v in self.__dict__.items():
            setattr(result, k, deepcopy(v))
        result.id = self.__generate_id()
        return result

class B(A):
   def do_b(self):
       pass

I currently replaced ??? with 'A'. This works, but if I want to clone an object of type B, I have to cast the return value if I want to call B-specific methods on it:

b = B().clone()
b.do_b()  # type error!

b = cast(B, B().clone())
b.do_b()  # OK

However, it's guaranteed that calling .clone() on an object of type B will return another object of type B, and this smacks rather too much of Java for my taste. Is there some generic-based way I can tell mypy that the method returns an object of type __class__, whatever that class is?

回答1:

You can do this by using generic methods with a generic self -- basically, annotate your self variable as being generic:

from typing import TypeVar

T = TypeVar('T', bound='A')

class A:
    def __generate_id(self) -> int:
        return 0

    def clone(self: T) -> T: 
        cls = self.__class__
        result = cls.__new__(cls)
        for k, v in self.__dict__.items():
            setattr(result, k, deepcopy(v))
        result.id = self.__generate_id()
        return result

class B(A):
    def do_b(self):
        pass

reveal_type(A().clone())  # Revealed type is A
reveal_type(B().clone())  # Revealed type is B

Basically, when we call clone(), the T typevar will be bound to the type of the current instance when mypy tries type-checking the call to clone. This ends up making the return type match the type of the instance, whatever it is.

You may be wondering why I set the upper bound of T to A. This is because of the line self.__generate_id(). Basically, in order for that line to type-check, self can't be literally any type: it needs to be A or some subclass of A. The bound encodes this requirement.

来源：https://stackoverflow.com/questions/57446257/can-mypy-select-a-method-return-type-based-on-the-type-of-the-current-object