How to specify return type for a function returning a pydantic type according to an arguement value?-CodePudding

I have two pydantic models, the second inheriting from the first:

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


class RandomSpecializedModel(RandomBaseModel):
    foo: str

Then I have a function that accepts some data and a model to use for instanciating a response:

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

Finally the result of this function is stored into typed variable:

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

This executes without any problem and works as expected, the do_something function instanciates a RandomBaseModel when response_model is omitted and instanciates a RandomSpecializedModel when it is instructed to use it. here is the output:

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

BUT this does not please mypy which complains with this message on line value_2: RandomSpecializedModel = do_something("42", response_model=RandomSpecializedModel):

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

How could I inform mypy that this function returns an instance of the pydantic model passed as the response_model argument?

To be clear: I am looking for a way to instruct mypy that this function could return a RandomBaseModel instance, a RandomSpecializedModel or an instance of any RandomBaseModel's subclass.

I found some similar question whose answers suggested to use a TypeVar, so I tried to change the do_something function for this:

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

Although it still executes as expected, mypy now complains with:

error: Incompatible default for argument "response_model" (default has type "Type[RandomBaseModel]", argument has type "Type[AnyRandomBaseModel]")

CodePudding user response：

You should probably make the base model generic instead.

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

CodePudding user response：

Try with an Union of types:

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

I know it's weird as one of your class inherits from the other but I think that you do not have the choice