如何不使用cast让mypy自动推断Widget子类的实际类型

以下是三种不同场景的优化实现方式，均可以避免重复编写cast语句：

1. 泛型化TestCase基类（最优方案，静态类型安全） #

你可以将TestCase定义为泛型基类，让每个子类自定义WIDGETS字典的值类型，mypy会自动推导每个子类的WIDGETS实际类型，完全不需要强制转换：

from dataclasses import dataclass
from typing import ClassVar, Dict, TypeVar, Generic

@dataclass
class Widget:
    """Generic class for widget"""

@dataclass
class Rectangle(Widget):
    """A Color Rectangle"""
    posx: int
    posy: int
    width: int = 500
    height: int = 200
    color: int = 0xFF0000

@dataclass
class Button(Widget):
    """A clickable button"""
    posx: int
    posy: int
    width: int = 200
    height: int = 100
    label: str = "some label"

# 定义泛型参数
WidgetT = TypeVar('WidgetT', bound=Widget)
@dataclass
class TestCase(Generic[WidgetT]):
    ID: ClassVar[int]
    WIDGETS: ClassVar[Dict[str, WidgetT]]

# 子类继承时指定WIDGETS的值类型
class TestCaseButton(TestCase[Button]):
    ID = 1
    WIDGETS = {
        "ref": Button(posx=0, posy=0, label="Click me!"),
    }

# 如果字典混合多种Widget类型，用Union声明即可
class TestCaseRectangle(TestCase[Button | Rectangle]):
    ID = 2
    WIDGETS = {
        "ref": Rectangle(posx=0, posy=0, color=0xFFFFFF),
        "other": Button(posx=100, posy=100, label="DANGER!"),
    }

# 函数参数直接指定具体的TestCase子类类型
def run_tc_button(tc: TestCaseButton):
    w = tc.WIDGETS["ref"]
    print(w.label) # mypy自动识别为Button类型，无报错

def run_tc_rectangle(tc: TestCaseRectangle):
    w = tc.WIDGETS["ref"]
    # 混合类型场景加一层isinstance判断即可窄化类型
    if isinstance(w, Rectangle):
        print(w.color) # 无类型报错

2. 封装泛型获取辅助函数（兼容原有代码结构） #

如果你不想修改原有基类的定义，可以封装一个通用的获取Widget的辅助函数，把cast逻辑收拢在一处，调用时无需重复写cast：

from typing import Type, TypeVar, cast
T = TypeVar('T', bound=Widget)

def get_widget(widgets: Dict[str, Widget], key: str, expected_type: Type[T]) -> T:
    return cast(T, widgets[key])

# 调用方式
def run_tc_button(tc: TestCase):
    w = get_widget(tc.WIDGETS, "ref", Button)
    print(w.label)

def run_tc_rectangle(tc: TestCase):
    w = get_widget(tc.WIDGETS, "ref", Rectangle)
    print(w.color)

3. 运行时类型窄化（更安全，兼顾运行时校验） #

如果需要兼顾运行时类型合法性，可以用isinstance做类型判断，mypy会自动窄化变量类型，不需要任何cast，同时可以捕获运行时的类型不匹配问题：

def run_tc_button(tc: TestCase):
    w = tc.WIDGETS["ref"]
    if isinstance(w, Button):
        print(w.label)
    else:
        raise TypeError(f"期望得到Button类型，实际得到{type(w)}")

def run_tc_rectangle(tc: TestCase):
    w = tc.WIDGETS["ref"]
    if isinstance(w, Rectangle):
        print(w.color)
    else:
        raise TypeError(f"期望得到Rectangle类型，实际得到{type(w)}")

内容的提问来源于stack exchange，提问作者Jean-Francois T.