@@ -162,6 +162,8 @@ class UnionProperty(Property):
162162 inner_properties : List [Property ]
163163 template : ClassVar [str ] = "union_property.pyi"
164164 has_properties_without_templates : bool = attr .ib (init = False )
165+ discriminator_property : Optional [str ]
166+ discriminator_mappings : Dict [str , Property ]
165167
166168 def __attrs_post_init__ (self ) -> None :
167169 super ().__attrs_post_init__ ()
@@ -171,6 +173,8 @@ def __attrs_post_init__(self) -> None:
171173
172174 def _get_inner_type_strings (self , json : bool = False ) -> List [str ]:
173175 inner_types = [p .get_type_string (no_optional = True , json = json ) for p in self .inner_properties ]
176+ if not json :
177+ inner_types .append ("UnknownType" )
174178 unique_inner_types = list (dict .fromkeys (inner_types ))
175179 return unique_inner_types
176180
@@ -428,13 +432,27 @@ def build_union_property(
428432 * , data : oai .Schema , name : str , required : bool , schemas : Schemas , parent_name : str
429433) -> Tuple [Union [UnionProperty , PropertyError ], Schemas ]:
430434 sub_properties : List [Property ] = []
435+ inverted_mappings = {}
436+ for k , v in (data .discriminator .mapping if data .discriminator else {}).items ():
437+ class_name = Reference .from_ref (v ).class_name
438+ if class_name in inverted_mappings :
439+ raise ArgumentError (
440+ f"Mapping more than one name to a class is currently not supported (class: { class_name } )."
441+ )
442+ inverted_mappings [Reference .from_ref (v ).class_name ] = k
443+ discriminator_mappings : Dict [str , Property ] = {}
431444 for sub_prop_data in chain (data .anyOf , data .oneOf ):
432445 sub_prop , schemas = property_from_data (
433446 name = name , required = required , data = sub_prop_data , schemas = schemas , parent_name = parent_name
434447 )
435448 if isinstance (sub_prop , PropertyError ):
436449 return PropertyError (detail = f"Invalid property in union { name } " , data = sub_prop_data ), schemas
450+
437451 sub_properties .append (sub_prop )
452+ if data .discriminator is not None :
453+ discriminated_by = inverted_mappings .get (sub_prop .reference .class_name )
454+ if discriminated_by is not None :
455+ discriminator_mappings [discriminated_by ] = sub_prop
438456
439457 default = convert_chain ((prop ._type_string for prop in sub_properties ), data .default )
440458 return (
@@ -444,6 +462,8 @@ def build_union_property(
444462 default = default ,
445463 inner_properties = sub_properties ,
446464 nullable = data .nullable ,
465+ discriminator_property = data .discriminator .propertyName if data .discriminator else None ,
466+ discriminator_mappings = discriminator_mappings ,
447467 description = data .description ,
448468 ),
449469 schemas ,
0 commit comments