如何根据参数值指定返回 pydantic 类型的函数的返回类型？

Question

我有两个 pydantic 模型，第二个继承第一个：

class RandomBaseModel(pydantic.BaseModel):
    foo: typing.Any


class RandomSpecializedModel(RandomBaseModel):
    foo: str

然后我有一个接受一些数据的函数和一个用于实例化响应的模型：

def do_something(
    data: typing.Any,
    response_model: typing.Type[RandomBaseModel] = RandomBaseModel
) -> RandomBaseModel:
    response = response_model(foo=data)
    print(f"---{type(response)}---")
    return response

最后这个函数的结果被存储到类型变量中：

value_1: RandomBaseModel = do_something(42)
value_2: RandomSpecializedModel = do_something("42", response_model=RandomSpecializedModel)

这执行时没有任何问题并按预期工作，

do_something

函数在省略

RandomBaseModel

时实例化

response_model

，并在指示使用它时实例化

RandomSpecializedModel

。这是输出：

---<class '__main__.RandomBaseModel'>---
---<class '__main__.RandomSpecializedModel'>---

但这并不能让 mypy 满意，它在网上抱怨此消息

value_2: RandomSpecializedModel = do_something("42", response_model=RandomSpecializedModel)

：

error: Incompatible types in assignment (expression has type "RandomBaseModel", variable has type "RandomSpecializedModel")  [assignment]

我如何通知 mypy 该函数返回作为

response_model

参数传递的 pydantic 模型的实例？

需要明确的是：我正在寻找一种方法来指示 mypy 该函数可以返回

RandomBaseModel

实例、

RandomSpecializedModel

或任何

RandomBaseModel

子类的实例。

我发现了一些类似的问题，其答案建议使用 TypeVar，所以我尝试为此更改

do_something

函数：

AnyRandomBaseModel = typing.TypeVar("AnyRandomBaseModel", bound=RandomBaseModel)


def do_something(
    data: typing.Any,
    response_model: typing.Type[AnyRandomBaseModel] = RandomBaseModel
) -> AnyRandomBaseModel:
    response = response_model(foo=data)
    print(f"---{type(response)}---")
    return response

虽然它仍然按预期执行，但 mypy 现在抱怨：

错误：参数“response_model”的默认值不兼容（默认值的类型为“Type[RandomBaseModel]”，参数的类型为“Type[AnyRandomBaseModel]”）

Answer 1

您可能应该将基本模型设为通用。

from typing import TypeVar, Generic

T = TypeVar('T')


class RandomBaseModel(pydantic.BaseModel, Generic[T]):
    foo: T


class RandomSpecializedModel(RandomBaseModel[str]):
    foo: str  # you might not need this line

def do_something(
    data: T,
    response_model: typing.Type[RandomBaseModel[T]] = RandomBaseModel
) -> RandomBaseModel[T]:
    response = response_model(foo=data)
    print(f"---{type(response)}---")
    return response

Answer 2

尝试使用类型联合：

def do_something(
    data: typing.Any,
    response_model: typing.Type[RandomBaseModel] = RandomBaseModel
) -> Union[RandomBaseModel,RandomSpecializedModel]

我知道这很奇怪，因为你的一个类继承了另一个类，但我认为你没有选择

Answer 3

仅供参考，因为我对此不太满意，所以我用

typing.cast

解决了我的问题：

value_2: RandomSpecializedModel = typing.cast(
    RandomSpecializedModel,
    do_something("42", response_model=RandomSpecializedModel)
)

为了取悦 Mypy 而使其如此冗长，这对我来说并不是很令人满意，但至少它可以在不使用

# type: ignore[]

静音 Mypy 的情况下工作。 所以这已经足够好了...:/