@@ -162,6 +162,8 @@ class UnionProperty(Property):
162
162
inner_properties : List [Property ]
163
163
template : ClassVar [str ] = "union_property.pyi"
164
164
has_properties_without_templates : bool = attr .ib (init = False )
165
+ discriminator_property : Optional [str ]
166
+ discriminator_mappings : Dict [str , Property ]
165
167
166
168
def __attrs_post_init__ (self ) -> None :
167
169
super ().__attrs_post_init__ ()
@@ -171,6 +173,8 @@ def __attrs_post_init__(self) -> None:
171
173
172
174
def _get_inner_type_strings (self , json : bool = False ) -> List [str ]:
173
175
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" )
174
178
unique_inner_types = list (dict .fromkeys (inner_types ))
175
179
return unique_inner_types
176
180
@@ -428,13 +432,27 @@ def build_union_property(
428
432
* , data : oai .Schema , name : str , required : bool , schemas : Schemas , parent_name : str
429
433
) -> Tuple [Union [UnionProperty , PropertyError ], Schemas ]:
430
434
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 ] = {}
431
444
for sub_prop_data in chain (data .anyOf , data .oneOf ):
432
445
sub_prop , schemas = property_from_data (
433
446
name = name , required = required , data = sub_prop_data , schemas = schemas , parent_name = parent_name
434
447
)
435
448
if isinstance (sub_prop , PropertyError ):
436
449
return PropertyError (detail = f"Invalid property in union { name } " , data = sub_prop_data ), schemas
450
+
437
451
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
438
456
439
457
default = convert_chain ((prop ._type_string for prop in sub_properties ), data .default )
440
458
return (
@@ -444,6 +462,8 @@ def build_union_property(
444
462
default = default ,
445
463
inner_properties = sub_properties ,
446
464
nullable = data .nullable ,
465
+ discriminator_property = data .discriminator .propertyName if data .discriminator else None ,
466
+ discriminator_mappings = discriminator_mappings ,
447
467
description = data .description ,
448
468
),
449
469
schemas ,
0 commit comments