66`register_schema_models` emit resolvable OpenAPI 3 references.
77"""
88
9+ import re
910from collections .abc import Iterable , Mapping
1011from enum import StrEnum
11- from typing import Any , Literal , NotRequired , Protocol , TypedDict
12+ from typing import Any , Literal , NotRequired , Protocol , TypedDict , override
13+ from weakref import WeakKeyDictionary
1214
1315from flask import request
1416from flask_restx import Namespace
1517from pydantic import BaseModel , TypeAdapter
18+ from pydantic .json_schema import CoreModeRef , DefsRef , GenerateJsonSchema
1619
1720DEFAULT_REF_TEMPLATE_OPENAPI_3_0 = "#/components/schemas/{model}"
1821
@@ -51,39 +54,125 @@ def to_dict(self, flat: bool = True) -> dict[str, str]: ...
5154 def getlist (self , key : str ) -> list [str ]: ...
5255
5356
54- def _register_json_schema (namespace : Namespace , name : str , schema : dict ) -> None :
55- """Register a JSON schema and promote any nested Pydantic `$defs`."""
57+ class SchemaNameCollisionError (ValueError ):
58+ """Raised when two different schemas would be published under the same OpenAPI component name."""
59+
60+
61+ # Pydantic names `$defs` after the bare class name and only qualifies them when two
62+ # same-named types meet inside one schema, so the same type can surface under different
63+ # names depending on which model pulled it in. Types that share a class name with
64+ # another documented type get a fixed, unique component name here instead.
65+ SCHEMA_NAME_OVERRIDES : Mapping [str , str ] = {
66+ "graphon.model_runtime.entities.common_entities.I18nObject" : "ModelRuntimeI18nObject" ,
67+ "services.knowledge.entities.knowledge_entities.FileInfo" : "KnowledgeFileInfo" ,
68+ "core.plugin.entities.plugin.PluginDependency.Github" : "PluginDependencyGithub" ,
69+ "core.plugin.entities.plugin.PluginDependency.Marketplace" : "PluginDependencyMarketplace" ,
70+ "core.plugin.entities.plugin.PluginDependency.Package" : "PluginDependencyPackage" ,
71+ "core.plugin.entities.bundle.PluginBundleDependency.Github" : "PluginBundleDependencyGithub" ,
72+ "core.plugin.entities.bundle.PluginBundleDependency.Marketplace" : "PluginBundleDependencyMarketplace" ,
73+ "core.plugin.entities.bundle.PluginBundleDependency.Package" : "PluginBundleDependencyPackage" ,
74+ "core.app.entities.task_entities.WorkflowAppBlockingResponse.Data" : "WorkflowAppBlockingResponseData" ,
75+ }
76+
77+ _REF_ID_SUFFIX = re .compile (r":\d+(?=$|[\[\],])" )
78+ _MODULE_PREFIX = re .compile (r"(?:[^.[\],]+\.)+([^.[\],]+)" )
79+
80+ # Normalized component name -> (component name, fully qualified source type, schema), per namespace.
81+ _registered_schemas : WeakKeyDictionary [Namespace , dict [str , tuple [str , str , dict ]]] = WeakKeyDictionary ()
82+ _NON_ALPHANUMERIC = re .compile (r"[^a-z0-9]" )
83+
84+
85+ def _normalized_schema_name (name : str ) -> str :
86+ """Fold names that client code generators turn into the same identifier (`JSONValue` / `JsonValue`)."""
87+
88+ return _NON_ALPHANUMERIC .sub ("" , name .lower ())
89+
90+
91+ class _OpenApiJsonSchemaGenerator (GenerateJsonSchema ):
92+ """Emit stable component names and remember which type produced each of them."""
93+
94+ def __init__ (self , * args : Any , ** kwargs : Any ) -> None :
95+ super ().__init__ (* args , ** kwargs )
96+ self .schema_sources : dict [str , str ] = {}
97+
98+ @override
99+ def get_defs_ref (self , core_mode_ref : CoreModeRef ) -> DefsRef :
100+ core_ref , _mode = core_mode_ref
101+ source = _REF_ID_SUFFIX .sub ("" , core_ref )
102+ name = SCHEMA_NAME_OVERRIDES .get (source ) or self .normalize_name (_MODULE_PREFIX .sub (r"\1" , source ))
103+ existing_source = self .schema_sources .setdefault (name , source )
104+ if existing_source != source :
105+ raise SchemaNameCollisionError (
106+ f"OpenAPI schema name { name !r} is produced by both { existing_source } and { source } ; "
107+ "rename one of them or add it to SCHEMA_NAME_OVERRIDES."
108+ )
109+ defs_ref = DefsRef (name )
110+ self ._prioritized_defsref_choices [defs_ref ] = [defs_ref ]
111+ return defs_ref
112+
113+
114+ def _register_json_schema (namespace : Namespace , name : str , schema : dict , sources : Mapping [str , str ]) -> None :
115+ """Register a JSON schema and promote any nested Pydantic `$defs`.
116+
117+ `sources` maps each component name to the type that produced it. Flask-RESTX keeps
118+ one flat model namespace, so a second registration under an existing name would
119+ silently replace the first. Re-registering the same type with the same schema is a
120+ no-op; a different type, or the same type with a different shape (for example a
121+ response model also registered in validation mode), raises instead.
122+ """
56123
57124 nested_definitions = schema .get ("$defs" )
58125 schema_to_register = dict (schema )
59126 if isinstance (nested_definitions , dict ):
60127 schema_to_register .pop ("$defs" )
61128
129+ source = sources .get (name , name )
130+ registered = _registered_schemas .setdefault (namespace , {})
131+ normalized_name = _normalized_schema_name (name )
132+ existing = registered .get (normalized_name )
133+ if existing is not None and existing != (name , source , schema_to_register ):
134+ existing_name , existing_source , _ = existing
135+ if existing_name != name :
136+ reason = f"collides with { existing_name !r} from { existing_source } once normalized by client generators"
137+ elif existing_source != source :
138+ reason = f"is produced by both { existing_source } and { source } "
139+ else :
140+ reason = f"is registered by { source } with two different shapes (validation vs serialization mode?)"
141+ raise SchemaNameCollisionError (
142+ f"OpenAPI schema name { name !r} in namespace { namespace .name !r} { reason } ; give each shape its own "
143+ "class name, add it to SCHEMA_NAME_OVERRIDES, or register response models only with "
144+ "register_response_schema_models."
145+ )
146+ registered [normalized_name ] = (name , source , schema_to_register )
147+
62148 namespace .schema_model (name , schema_to_register )
63149
64150 if not isinstance (nested_definitions , dict ):
65151 return
66152
67153 for nested_name , nested_schema in nested_definitions .items ():
68154 if isinstance (nested_schema , dict ):
69- _register_json_schema (namespace , nested_name , nested_schema )
155+ _register_json_schema (namespace , nested_name , nested_schema , sources )
70156
71157
72158JsonSchemaMode = Literal ["validation" , "serialization" ]
73159
74160
75- def _register_schema_model (namespace : Namespace , model : type [BaseModel ], * , mode : JsonSchemaMode ) -> None :
161+ def _register_type_schema (namespace : Namespace , source_type : type , * , mode : JsonSchemaMode ) -> None :
162+ generator = _OpenApiJsonSchemaGenerator (ref_template = DEFAULT_REF_TEMPLATE_OPENAPI_3_0 )
163+ schema = generator .generate (TypeAdapter (source_type ).core_schema , mode = mode )
76164 _register_json_schema (
77165 namespace ,
78- model .__name__ ,
79- model .model_json_schema (ref_template = DEFAULT_REF_TEMPLATE_OPENAPI_3_0 , mode = mode ),
166+ source_type .__name__ ,
167+ schema ,
168+ {** generator .schema_sources , source_type .__name__ : f"{ source_type .__module__ } .{ source_type .__qualname__ } " },
80169 )
81170
82171
83172def register_schema_model (namespace : Namespace , model : type [BaseModel ]) -> None :
84173 """Register a BaseModel and its nested component schemas for OpenAPI documentation."""
85174
86- _register_schema_model (namespace , model , mode = "validation" )
175+ _register_type_schema (namespace , model , mode = "validation" )
87176
88177
89178def register_schema_models (namespace : Namespace , * models : type [BaseModel ]) -> None :
@@ -96,7 +185,7 @@ def register_schema_models(namespace: Namespace, *models: type[BaseModel]) -> No
96185def register_response_schema_model (namespace : Namespace , model : type [BaseModel ]) -> None :
97186 """Register a BaseModel using its serialized response shape."""
98187
99- _register_schema_model (namespace , model , mode = "serialization" )
188+ _register_type_schema (namespace , model , mode = "serialization" )
100189
101190
102191def register_response_schema_models (namespace : Namespace , * models : type [BaseModel ]) -> None :
@@ -119,11 +208,7 @@ def get_or_create_model(model_name: str, field_def):
119208def register_enum_models (namespace : Namespace , * models : type [StrEnum ]) -> None :
120209 """Register multiple StrEnum with a namespace."""
121210 for model in models :
122- _register_json_schema (
123- namespace ,
124- model .__name__ ,
125- TypeAdapter (model ).json_schema (ref_template = DEFAULT_REF_TEMPLATE_OPENAPI_3_0 ),
126- )
211+ _register_type_schema (namespace , model , mode = "validation" )
127212
128213
129214def query_params_from_model (model : type [BaseModel ]) -> dict [str , QueryParamDoc ]:
@@ -334,6 +419,8 @@ def _nullable_property_schema(property_schema: Mapping[str, Any]) -> Mapping[str
334419
335420__all__ = [
336421 "DEFAULT_REF_TEMPLATE_OPENAPI_3_0" ,
422+ "SCHEMA_NAME_OVERRIDES" ,
423+ "SchemaNameCollisionError" ,
337424 "get_or_create_model" ,
338425 "query_params_from_model" ,
339426 "query_params_from_request" ,
0 commit comments