如何为接受TypedDict联合与Literal联合的函数添加类型提示？

问题场景 #

我有一组TypedDict类型，需要从中提取信息。每个TypedDict对应一组允许传入函数的Literal值，且这些值存在重叠。是否有办法不用为每个TypedDict编写单独重载，就能给函数添加正确的类型提示？

测试示例代码 #

from typing import Literal, overload, TypedDict, LiteralString

class Test1(TypedDict):
    test1: str

class Test2(TypedDict):
    test2: str

@overload
def _get_actor_key(
    actor: Literal["test1"],
    game_action: Test1,
) -> str:
    ...

@overload
def _get_actor_key(
    actor: Literal["test2"], game_action: Test2
) -> str:
    ...

def _get_actor_key(
    actor: Literal["test1","test2"],
    game_action: Test1 | Test2,
) -> str:
    return game_action[actor]

实际业务场景示例 #

@overload
def _get_actor_key(
    actor: Literal["attacker", "victim", "assister", "flashThrower", "playerTraded"],
    game_action: KillAction,
) -> str:
    ...

@overload
def _get_actor_key(
    actor: Literal["attacker", "player"], game_action: FlashAction
) -> str:
    ...

def _get_actor_key(
    actor: GameActionPlayers,
    game_action: GameAction,
) -> str:  # type: ignore[reportGeneralTypeIssues]
    return (
        str(game_action[actor + "Name"])  # type: ignore[reportGeneralTypeIssues]
        if game_action[actor + "SteamID"] == 0  # type: ignore[reportGeneralTypeIssues]
        else str(game_action[actor + "SteamID"])  # type: ignore[reportGeneralTypeIssues]
    )

注：部分游戏动作支持获取attacker、victim等角色的名称和SteamID，其他动作仅支持获取player的相关信息。

解决方案 #

可以通过泛型+键类型约束替代重复的重载定义，利用typing.KeyOf（Python 3.11+内置，低版本可使用typing_extensions.KeyOf）直接关联TypedDict的键与函数参数的Literal类型，无需为每个TypedDict编写单独重载。

针对测试场景的简化实现 #

from typing import Generic, TypeVar, TypedDict
from typing_extensions import KeyOf  # Python<3.11时需安装typing-extensions

# 定义绑定到TypedDict的类型变量
T = TypeVar("T", bound=TypedDict)

def _get_actor_key(actor: KeyOf[T], game_action: T) -> str:
    return game_action[actor]

此时类型检查器会自动推断：传入Test1实例时，actor只能是"test1"；传入Test2实例时，actor只能是"test2"，完全替代原有的重载逻辑。

针对业务场景的优化实现 #

业务场景中，每个GameAction对应的actor需要关联{actor}Name和{actor}SteamID两个键，可以通过泛型协议约束关联关系，同时避免类型忽略注释：

from typing import Protocol, TypeVar, TypedDict, Literal
from typing_extensions import KeyOf

# 定义actor类型的类型变量
ActorType = TypeVar("ActorType", bound=str)

# 定义协议，约束GameAction必须包含actor对应的Name和SteamID键
class ActorActionProtocol(Protocol[ActorType]):
    def __getitem__(self, key: Literal[f"{ActorType}Name", f"{ActorType}SteamID"]) -> str | int: ...

# 定义各动作允许的actor类型
KillActor = Literal["attacker", "victim", "assister", "flashThrower", "playerTraded"]
FlashActor = Literal["attacker", "player"]

# 让各TypedDict实现协议
class KillAction(TypedDict, ActorActionProtocol[KillActor]):
    attackerName: str
    attackerSteamID: int
    victimName: str
    victimSteamID: int
    assisterName: str
    assisterSteamID: int
    flashThrowerName: str
    flashThrowerSteamID: int
    playerTradedName: str
    playerTradedSteamID: int

class FlashAction(TypedDict, ActorActionProtocol[FlashActor]):
    attackerName: str
    attackerSteamID: int
    playerName: str
    playerSteamID: int

# 定义泛型函数
T = TypeVar("T", bound=ActorActionProtocol)

def _get_actor_key(actor: T.__args__[0], game_action: T) -> str:
    name_key = f"{actor}Name"
    steam_id_key = f"{actor}SteamID"
    return str(game_action[name_key]) if game_action[steam_id_key] == 0 else str(game_action[steam_id_key])

实现后，类型检查器会自动校验：

• 传入KillAction时，actor只能是KillActor中定义的字面量
• 传入FlashAction时，actor只能是FlashActor中定义的字面量
• 函数体内的键访问不会再触发类型报错，无需添加type: ignore注释

内容的提问来源于stack exchange，提问作者J.N.