如何根据参数值指定返回类型的函数返回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)

这个执行没有任何问题，并按预期工作，当response_model被省略时，do_something函数实例化一个RandomBaseModel，当它被指示使用它时，实例化一个RandomSpecializedModel。下面是输出:

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

但这并不让我满意，我在value_2: RandomSpecializedModel = do_something("42", response_model=RandomSpecializedModel)线上抱怨这条消息:

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

我怎么能告诉我，这个函数返回pydantic模型的实例传递作为response_model参数?

---

要明确:我正在寻找一种方法来指示我的函数可以返回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现在报错:

error: Incompatible default for argument "response_model" (default has type "Type[RandomBaseModel]", argument has type "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

尝试使用类型的Union:

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的情况下工作。所以这就足够好了……:/