Python 根据属性类型定义返回类型

Question

我有三个类，一个基类和两个具有以下结构的子类：

class BaseLayout:
    def add_element(self: "BaseLayout", element: str) -> None:
        ...

class LayoutA(BaseLayout):
    def add_element(self: "BaseLayout", element: str, name: str) -> None:
        ...

class LayoutB(BaseLayout):
    def add_element(self: "BaseLayout", element: str, number: int) -> None:
        ...

基类没有直接使用，但确实实现了一些与问题无关的其他功能。 然后，我将布局用作 Container 类中的属性：

class Container:
    def __init__(self: "Container", layout: BaseLayout) -> None:
        self._current_layout = layout

    @property
    def layout(self: "Container") -> BaseLayout:
        return self._current_layout

这是我目前的问题：

如果我向Container类提供LayoutA的一个实例，然后调用add_element方法，我会得到类型检查器错误Expected 1 positional argument 。 下面是添加元素的代码：

c = Container(LayoutA())

# Expected 1 positional argument
c.add_element("foo", "bar")

我怎样才能向 Container 类提供正确的类型提示以使其工作？

Answer 1

正如@chepner 在评论中已经解释的那样，您的代码首先不是类型安全的，因为您的BaseLayout子类的add_element覆盖与基本签名不兼容。

例如，如果您需要不同的行为，您可以通过创建不同的方法来解决这个问题，并让它们调用超类的方法。 我不知道实际的用例，所以这只是一个例子。 有很多方法可以解决这个问题。

class BaseLayout:
    def add_element(self, element: str) -> None:
        print(f"{element=}")


class LayoutA(BaseLayout):
    def add_name_element(self, element: str, name: str) -> None:
        super().add_element(element)
        print(f"{name=}")


class LayoutB(BaseLayout):
    def add_number_element(self, element: str, number: int) -> None:
        super().add_element(element)
        print(f"{number=}")

修复后，您可以根据 Container 使用的布局使Container通用，如下所示：

from typing import Generic, TypeVar


L = TypeVar("L", bound=BaseLayout)


class Container(Generic[L]):
    _current_layout: L  # this is optional, but helpful IMO

    def __init__(self, layout: L) -> None:
        self._current_layout = layout

    @property
    def layout(self) -> L:
        return self._current_layout


ca = Container(LayoutA())
ca.layout.add_name_element("foo", "bar")

cb = Container(LayoutB())
cb.layout.add_number_element("foo", 1)

这很好，静态类型检查器应该正确推断layout类型。 如果你在下面添加reveal_type(ca._current_layout)和reveal_type(cb._current_layout)并运行mypy ，你会得到这个：

note: Revealed type is "LayoutA"
note: Revealed type is "LayoutB"

这也意味着，如果您尝试在某处调用cb.layout.add_name_element ，类型检查器会报错。