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 ，類型檢查器會報錯。