定义子类的返回值类型，而无需在 python 中重新定义父实现

Question

我有一个树的基本实现，它可以包含树和具有唯一 ID 的通用元素，并且我有一些树的子类，其中包含特定类型的元素，如下所示：

class BaseTree:

    def __init__(self):
        self._my_elements_and_trees = []

    def append_element_or_tree(self, element_or_tree: Union["BaseTree", Any]) -> int:
        ...

    def iter_all_elements(self, is_recursive=True) -> Iterable[Any]:
        ...

    def iter_all_trees(self, is_recursive=True) -> Iterable["BaseTree"]:
        ...

    def get_element_from_id(self, element_id) -> Any:
        ...


class TreeOfMessages(BaseTree):
    """
    Tree containing instances of Messages
    """
    def do_something_with_messages(self):
        ...


class TreeOfMails(BaseTree):
    """
    Tree containing instances of Mails
    """

    def do_something_with_mails(self):
        ...

我会声明我的类型，这样当我在我的任何特定子类上使用任何父方法（ iter_all_elements 、 iter_all_trees或get_element_from_id ）时，我将获得特定对象（消息或邮件）及其相应的树。

我知道可以执行以下操作：

class TreeOfMessages(BaseTree):
    """
    Tree containing instances of Messages
    """

    def get_element_from_id(self, element_id) -> Message:
        return super().get_element_from_id(element_id)

但我不想每次都重新实现每个方法只是为了声明类型。

Answer 1

您正在寻找的是Generics 。

您可以将 BaseTree 实现为具有特定类型的通用类，如下所示：

from typing import Generic, TypeVar

ContainedType = TypeVar("ContainedType")

class BaseTree(Generic[ContainedType]):

    # Now just replace all parts where your current implementation uses "Any"
    # with "ContainedType", e.g.
    def get_element_from_id(self, element_id) -> ContainedType:
        ...

然后你可以让你的子类承认它只能通过这样定义它来接收消息：

class MyMessage:
    pass

class TreeOfMessages(BaseTree[MyMessage]):
    pass

并且类型检查器将推断 TreeOfMessages 使用 MyMessage 作为所有调用的包含类型。