heterogeneous.py

from pydantic import BaseModel, Field, RootModel, Discriminator, Tag
from typing import Literal, TypeVar, Generic, Annotated, Any

_ResponseType = TypeVar("_ResponseType")


def get_discriminator_value(v: Any) -> str:
    if isinstance(v, dict):
        return "success" if v.get("code") == 0 else "error"
    return "success" if getattr(v, "code") == 0 else "error"


class _SuccessResponse(BaseModel, Generic[_ResponseType]):
    code: Literal[0] = 0
    data: _ResponseType


class _ErrorResponse(BaseModel):
    code: int
    message: str
    suggestion: str


class BaseResponse(RootModel[_ResponseType]):
    root: Annotated[
        Annotated[_SuccessResponse[_ResponseType], Tag("success")]
        | Annotated[_ErrorResponse, Tag("error")],
        Field(discriminator=Discriminator(get_discriminator_value)),
    ]

    def __getattr__(self, item):
        return getattr(self.root, item)

    @property
    def success(self) -> bool:
        return self.root.code == 0


class Test(BaseModel):
    a: int
    b: str


if __name__ == "__main__":
    item = BaseResponse[Test].model_validate_json(
        '{"code": 0, "data": {"a": 1, "b": "2"}}'
    )

    print(type(item.data))