用于获取与静态类型检查器一起使用的 TypedDict 值类型的函数

Question

我正在寻找一个 Python (3.8+) 函数：

• 输入：TypedDict 的键
• 输出：
  • 返回值（简单）
  • 是否有适当的类型提示（我被卡住的地方）

这里有一个代码示例来帮助解释：

from typing import Any, Literal, TypedDict

class Foo(TypedDict):
    bar: int
    baz: str
    spam: Any

foo = Foo(bar=0, baz="hi", spam=1.0)

def get_type(key: Literal["bar", "baz"]):  # How to type hint the return here?
    """Function that get TypedDict's value when passed a key."""
    val = foo[key]
    # This works via intelligent indexing
    # SEE: https://mypy.readthedocs.io/en/stable/literal_types.html#intelligent-indexing
    reveal_type(val)  # mypy: Revealed type is 'Union[builtins.int, builtins.str]'
    return val

fetched_type = get_type("bar")
reveal_type(fetched_type)  # mypy: Revealed type is 'Any'
# I would like this to have output: 'int'

如果你看不出来，我使用的静态类型检查器是mypy。

我在get_type 上面的函数在intelligent indexing 中实现了一半，但我不知道如何键入提示get_type 的返回。

返回get_type应该输入什么类型的提示？

---

研究

这两个问题

• Factory function for mypy `TypedDict`
• How to statically get TypeVar parameters from a Generic for use in static type checking?

使用TypeVar 提供答案。有什么方法可以将TypeVar 与TypedDict 一起使用？

Answer 1

如果我正确理解您的问题，您可以将@overload 用于get_type：

from typing import Any, Literal, TypedDict, Union, overload

class Foo(TypedDict):
    bar: int
    baz: str
    spam: Any

foo = Foo(bar=0, baz="hi", spam=1.0)


@overload
def get_type(key: Literal["bar"]) -> int: ...

@overload
def get_type(key: Literal["baz"]) -> str: ...
    

def get_type(key: Literal["bar", "baz"]) -> Union[int, str]:
    """Function that get TypedDict's value when passed a key."""
    val = foo[key]
    reveal_type(val)  # mypy: Revealed type is 'Union[builtins.int, builtins.str]'
    return val

fetched_type = get_type("bar")
reveal_type(fetched_type)  # mypy: Revealed type is 'builtins.int'