@@ -1246,6 +1246,9 @@ private ModelsMap postProcessModelsMap(ModelsMap objs) {
1246
1246
pydanticImports .add ("validator" );
1247
1247
} else { // typical model
1248
1248
codegenProperties = model .vars ;
1249
+ if (model .getDiscriminator () != null && model .getDiscriminator ().getMappedModels () != null ) {
1250
+ typingImports .add ("Union" );
1251
+ }
1249
1252
}
1250
1253
1251
1254
//loop through properties/schemas to set up typing, pydantic
@@ -1323,39 +1326,37 @@ private ModelsMap postProcessModelsMap(ModelsMap objs) {
1323
1326
}
1324
1327
}
1325
1328
1326
- if (!model .isEnum ) {
1327
- pydanticImports .add ("BaseModel" );
1328
- }
1329
-
1330
1329
// add parent model to import
1331
1330
if (!StringUtils .isEmpty (model .parent )) {
1332
1331
modelImports .add (model .parent );
1333
1332
// Extention for discriminator
1334
1333
modelImports .addAll (model .imports );
1335
1334
1336
1335
String mappedTypeName = mappingName (model .name , model .parentModel .getDiscriminator ());
1337
- model .vendorExtensions .putIfAbsent ("x-py-type-name" , mappedTypeName );
1336
+ model .vendorExtensions .put ("x-py-type-name" , mappedTypeName );
1337
+ } else if (!model .isEnum ) {
1338
+ pydanticImports .add ("BaseModel" );
1338
1339
}
1339
1340
1340
1341
// set enum type in extensions and update `name` in enumVars
1341
1342
if (model .isEnum ) {
1342
1343
for (Map <String , Object > enumVars : (List <Map <String , Object >>) model .getAllowableValues ().get ("enumVars" )) {
1343
1344
if ((Boolean ) enumVars .get ("isString" )) {
1344
- model .vendorExtensions .putIfAbsent ("x-py-enum-type" , "str" );
1345
+ model .vendorExtensions .put ("x-py-enum-type" , "str" );
1345
1346
// update `name`, e.g.
1346
1347
enumVars .put ("name" , toEnumVariableName ((String ) enumVars .get ("value" ), "str" ));
1347
1348
} else {
1348
- model .vendorExtensions .putIfAbsent ("x-py-enum-type" , "int" );
1349
+ model .vendorExtensions .put ("x-py-enum-type" , "int" );
1349
1350
enumVars .put ("name" , toEnumVariableName ((String ) enumVars .get ("value" ), "int" ));
1350
1351
}
1351
1352
}
1352
1353
}
1353
1354
1354
1355
// set the extensions if the key is absent
1355
- model .getVendorExtensions ().putIfAbsent ("x-py-typing-imports" , typingImports );
1356
- model .getVendorExtensions ().putIfAbsent ("x-py-pydantic-imports" , pydanticImports );
1357
- model .getVendorExtensions ().putIfAbsent ("x-py-datetime-imports" , datetimeImports );
1358
- model .getVendorExtensions ().putIfAbsent ("x-py-readonly" , readOnlyFields );
1356
+ model .getVendorExtensions ().put ("x-py-typing-imports" , typingImports );
1357
+ model .getVendorExtensions ().put ("x-py-pydantic-imports" , pydanticImports );
1358
+ model .getVendorExtensions ().put ("x-py-datetime-imports" , datetimeImports );
1359
+ model .getVendorExtensions ().put ("x-py-readonly" , readOnlyFields );
1359
1360
1360
1361
// import models one by one
1361
1362
if (!modelImports .isEmpty ()) {
0 commit comments