163 lines
4.9 KiB
Python
163 lines
4.9 KiB
Python
from typing import Dict, List, Union
|
|
|
|
from pydantic import BaseModel, Field
|
|
from typing_extensions import Literal
|
|
|
|
from openapi_pydantic import (
|
|
Discriminator,
|
|
Info,
|
|
MediaType,
|
|
OpenAPI,
|
|
Operation,
|
|
PathItem,
|
|
Reference,
|
|
RequestBody,
|
|
Response,
|
|
Schema,
|
|
)
|
|
from openapi_pydantic.compat import (
|
|
DEFS_KEY,
|
|
PYDANTIC_MINOR_VERSION,
|
|
PYDANTIC_V2,
|
|
models_json_schema,
|
|
v1_schema,
|
|
)
|
|
from openapi_pydantic.util import PydanticSchema, construct_open_api_with_schema_class
|
|
|
|
|
|
class DataAModel(BaseModel):
|
|
kind: Literal["a"]
|
|
|
|
|
|
class DataBModel(BaseModel):
|
|
kind: Literal["b"]
|
|
|
|
|
|
class RequestModel(BaseModel):
|
|
data: Union[DataAModel, DataBModel] = Field(discriminator="kind")
|
|
|
|
|
|
def construct_base_open_api() -> OpenAPI:
|
|
return OpenAPI(
|
|
info=Info(
|
|
title="My own API",
|
|
version="v0.0.1",
|
|
),
|
|
paths={
|
|
"/ping": PathItem(
|
|
post=Operation(
|
|
requestBody=RequestBody(
|
|
content={
|
|
"application/json": MediaType(
|
|
media_type_schema=PydanticSchema(
|
|
schema_class=RequestModel
|
|
)
|
|
)
|
|
}
|
|
),
|
|
responses={"200": Response(description="pong")},
|
|
)
|
|
)
|
|
},
|
|
)
|
|
|
|
|
|
def test_pydantic_discriminator_schema_generation() -> None:
|
|
"""https://github.com/kuimono/openapi-schema-pydantic/issues/8"""
|
|
|
|
a_kind: Dict[str, Union[str, List[str]]] = {"title": "Kind", "type": "string"}
|
|
b_kind: Dict[str, Union[str, List[str]]] = {"title": "Kind", "type": "string"}
|
|
|
|
if PYDANTIC_V2:
|
|
_key_map, json_schema = models_json_schema([(RequestModel, "validation")])
|
|
# In Pydantic v2, string literal types are mapped to consts.
|
|
a_kind["const"] = "a"
|
|
b_kind["const"] = "b"
|
|
if PYDANTIC_MINOR_VERSION < 10:
|
|
# Prior to 2.10, string literal types are also mapped to enums with a single entry.
|
|
a_kind["enum"] = ["a"]
|
|
b_kind["enum"] = ["b"]
|
|
else:
|
|
json_schema = v1_schema([RequestModel])
|
|
# In Pydantic v1, string literal types are mapped to enums with a single entry.
|
|
a_kind["enum"] = ["a"]
|
|
b_kind["enum"] = ["b"]
|
|
|
|
assert json_schema == {
|
|
DEFS_KEY: {
|
|
"DataAModel": {
|
|
"properties": {
|
|
"kind": a_kind,
|
|
},
|
|
"required": ["kind"],
|
|
"title": "DataAModel",
|
|
"type": "object",
|
|
},
|
|
"DataBModel": {
|
|
"properties": {
|
|
"kind": b_kind,
|
|
},
|
|
"required": ["kind"],
|
|
"title": "DataBModel",
|
|
"type": "object",
|
|
},
|
|
"RequestModel": {
|
|
"properties": {
|
|
"data": {
|
|
"oneOf": [
|
|
{"$ref": f"#/{DEFS_KEY}/DataAModel"},
|
|
{"$ref": f"#/{DEFS_KEY}/DataBModel"},
|
|
],
|
|
"discriminator": {
|
|
"mapping": {
|
|
"a": f"#/{DEFS_KEY}/DataAModel",
|
|
"b": f"#/{DEFS_KEY}/DataBModel",
|
|
},
|
|
"propertyName": "kind",
|
|
},
|
|
"title": "Data",
|
|
}
|
|
},
|
|
"required": ["data"],
|
|
"title": "RequestModel",
|
|
"type": "object",
|
|
},
|
|
}
|
|
}
|
|
|
|
|
|
def test_pydantic_discriminator_openapi_generation() -> None:
|
|
"""https://github.com/kuimono/openapi-schema-pydantic/issues/8"""
|
|
|
|
open_api = construct_open_api_with_schema_class(construct_base_open_api())
|
|
assert open_api.components is not None
|
|
assert open_api.components.schemas is not None
|
|
json_schema = open_api.components.schemas["RequestModel"]
|
|
assert json_schema.properties == {
|
|
"data": Schema(
|
|
oneOf=[
|
|
Reference(
|
|
**{
|
|
"$ref": "#/components/schemas/DataAModel",
|
|
"summary": None,
|
|
"description": None,
|
|
}
|
|
),
|
|
Reference(
|
|
**{
|
|
"$ref": "#/components/schemas/DataBModel",
|
|
"summary": None,
|
|
"description": None,
|
|
}
|
|
),
|
|
],
|
|
title="Data",
|
|
discriminator=Discriminator(
|
|
propertyName="kind",
|
|
mapping={
|
|
"a": "#/components/schemas/DataAModel",
|
|
"b": "#/components/schemas/DataBModel",
|
|
},
|
|
),
|
|
)
|
|
}
|