如何为类似next()的函数实现动态返回类型注解？

Python内置next()函数的返回逻辑有特殊性：调用next(iterable)会返回迭代器元素，空迭代器则抛出StopIteration；调用next(iterable, default)则返回元素或指定默认值。若要自己实现同类逻辑的函数（比如示例中的MyDict.get()），如何添加类型注解让pyright这类工具正确识别不同调用场景的返回类型？

以下是最初的尝试代码，但无法通过类型检查：

from typing import TypeVar, Union

R = TypeVar('R')

class SuperNone:
    pass
    
class MyDict:
    data = { "a": 1 }

    def get(
        self, key: str, default: R = SuperNone
    ) -> Union[R, int] if R == SuperNone else int:
        try:
            return self.data[key]
        except KeyError:
            if isinstance(default, SuperNone):
                raise
            else:
                return default

a: int = MyDict().get("a")  # 报错："SuperNone | int" 无法赋值给类型 "int"
b: Union[int, str] = MyDict().get("a", "")

# 对比内置next()的正常表现：
c: int = next((x for x in [1]))
d: Union[int, str] = next((x for x in [1]), "")

解决方案：使用函数重载（@overload） #

要实现这种动态返回类型的注解，正确方式是用@overload定义多个函数签名，覆盖不同调用场景：

from typing import TypeVar, Union, overload

R = TypeVar('R')
Elem = int  # 明确字典值的类型

class MyDict:
    data = { "a": 1 }

    @overload
    def get(self, key: str) -> Elem:
        # 对应无默认参数的调用：仅返回Elem类型，找不到键则抛出异常
        ...

    @overload
    def get(self, key: str, default: R) -> Union[Elem, R]:
        # 对应带默认参数的调用：返回Elem或默认值类型R
        ...

    # 实际实现的函数体
    def get(self, key: str, default: Union[R, object] = object()) -> Union[Elem, R]:
        try:
            return self.data[key]
        except KeyError:
            if default is object():
                raise
            else:
                return default

# 现在类型检查可正常通过
a: int = MyDict().get("a")
b: Union[int, str] = MyDict().get("a", "")

关键细节 #

• 多签名覆盖场景：类型检查器会根据调用时的参数数量匹配对应重载签名，自动推断正确的返回类型。
• 占位符选择：用object()代替自定义的SuperNone作为默认参数占位符，无需额外定义类，且不会与用户传入的实际值冲突。
• 返回类型明确：无默认参数的场景直接返回元素类型，带默认参数的场景返回元素与默认值的联合类型，完全匹配函数的实际行为。

内容的提问来源于stack exchange，提问作者KamilCuk