feat: pydantic discriminator for unions (#2316)

Author: 412b

.sonarcloud.properties

@@ -1,3 +1,5 @@
 # Disable specific duplicate code since it would introduce more complexity to reduce it.
 sonar.cpd.exclusions=src/generators/**/*.ts
 sonar.exclusions=modelina-website/next.config.js
+sonar.sources=src,modelina-cli/src
+sonar.tests=test,modelina-cli/test

modelina-cli/src/helpers/python.ts

@@ -23,8 +23,8 @@ export function buildPythonGenerator(flags: any): BuilderReturnType {
   if (packageName === undefined) {
     throw new Error('In order to generate models to Python, we need to know which package they are under. Add `--packageName=PACKAGENAME` to set the desired package name.');
   }
-  
-  const presets = [];  
+
+  const presets = [];
   if (pyDantic) {
     presets.push(PYTHON_PYDANTIC_PRESET);
   }

src/generators/python/PythonConstrainer.ts

@@ -7,6 +7,11 @@ import { defaultModelNameConstraints } from './constrainer/ModelNameConstrainer'
 import { defaultPropertyKeyConstraints } from './constrainer/PropertyKeyConstrainer';
 import { defaultConstantConstraints } from './constrainer/ConstantConstrainer';
 import { PythonOptions, PythonTypeMapping } from './PythonGenerator';
+import { PYTHON_PYDANTIC_PRESET } from './presets';
+import {
+  ConstrainedObjectModel,
+  ConstrainedReferenceModel
+} from '../../models';
 
 export const PythonDefaultTypeMapping: PythonTypeMapping = {
   Object({ constrainedModel }): string {
@@ -46,12 +51,38 @@ export const PythonDefaultTypeMapping: PythonTypeMapping = {
     //Returning name here because all enum models have been split out
     return constrainedModel.name;
   },
-  Union({ constrainedModel }): string {
+  Union({ constrainedModel, options, dependencyManager }): string {
+    dependencyManager.addDependency('from typing import Union');
     const unionTypes = constrainedModel.union.map((unionModel) => {
       return unionModel.type;
     });
     const uniqueSet = new Set(unionTypes);
-    return [...uniqueSet].join(' | ');
+    // support Python pre 3.10
+    let union = `Union[${[...uniqueSet].join(', ')}]`;
+
+    if (
+      constrainedModel.options.discriminator &&
+      options.presets?.includes(PYTHON_PYDANTIC_PRESET)
+    ) {
+      const discriminator =
+        constrainedModel.options.discriminator.discriminator;
+      const reference = constrainedModel.union[0];
+      if (reference instanceof ConstrainedReferenceModel) {
+        const ref = reference.ref;
+        if (ref instanceof ConstrainedObjectModel) {
+          const properties = ref.properties;
+          const property = Object.values(properties).find((property) => {
+            return property.unconstrainedPropertyName === discriminator;
+          });
+          if (property !== undefined) {
+            dependencyManager.addDependency('from typing import Annotated');
+            union = `Annotated[${union}, Field(discriminator='${property.propertyName}')]`;
+          }
+        }
+      }
+    }
+
+    return union;
   },
   Dictionary({ constrainedModel }): string {
     return `dict[${constrainedModel.key.type}, ${constrainedModel.value.type}]`;

src/generators/python/presets/Pydantic.ts

@@ -1,15 +1,14 @@
 import {
   ConstrainedDictionaryModel,
-  ConstrainedObjectPropertyModel,
-  ConstrainedUnionModel
+  ConstrainedObjectPropertyModel
 } from '../../../models';
 import { PythonOptions } from '../PythonGenerator';
 import { ClassPresetType, PythonPreset } from '../PythonPreset';
 
 const PYTHON_PYDANTIC_CLASS_PRESET: ClassPresetType<PythonOptions> = {
   async self({ renderer, model }) {
     renderer.dependencyManager.addDependency(
-      'from typing import Optional, Any, Union'
+      'from typing import Optional, Any'
     );
     renderer.dependencyManager.addDependency(
       'from pydantic import BaseModel, Field'
@@ -26,18 +25,19 @@ const PYTHON_PYDANTIC_CLASS_PRESET: ClassPresetType<PythonOptions> = {
     let type = property.property.type;
     const propertyName = property.propertyName;
 
-    if (property.property instanceof ConstrainedUnionModel) {
-      const unionTypes = property.property.union.map(
-        (unionModel) => unionModel.type
-      );
-      type = `Union[${unionTypes.join(', ')}]`;
-    }
-
     const isOptional =
       !property.required || property.property.options.isNullable === true;
     if (isOptional) {
       type = `Optional[${type}]`;
     }
+    if (
+      property.property.options.const &&
+      model.options.discriminator?.discriminator ===
+        property.unconstrainedPropertyName
+    ) {
+      renderer.dependencyManager.addDependency('from typing import Literal');
+      type = `Literal['${property.property.options.const.originalInput}']`;
+    }
     type = renderer.renderPropertyType({
       modelType: model.type,
       propertyType: type
@@ -54,7 +54,14 @@ const PYTHON_PYDANTIC_CLASS_PRESET: ClassPresetType<PythonOptions> = {
       decoratorArgs.push('default=None');
     }
     if (property.property.options.const) {
-      decoratorArgs.push(`default=${property.property.options.const.value}`);
+      let value = property.property.options.const.value;
+      if (
+        model.options.discriminator?.discriminator ===
+        property.unconstrainedPropertyName
+      ) {
+        value = property.property.options.const.originalInput;
+      }
+      decoratorArgs.push(`default='${value}'`);
       decoratorArgs.push('frozen=True');
     }
     if (

test/generators/python/PythonConstrainer.spec.ts

@@ -196,7 +196,7 @@ describe('PythonConstrainer', () => {
         constrainedModel: model,
         ...defaultOptions
       });
-      expect(type).toEqual('str');
+      expect(type).toEqual('Union[str]');
     });
     test('should render multiple types', () => {
       const unionModel1 = new ConstrainedStringModel(
@@ -219,7 +219,7 @@ describe('PythonConstrainer', () => {
         constrainedModel: model,
         ...defaultOptions
       });
-      expect(type).toEqual('str | str2');
+      expect(type).toEqual('Union[str, str2]');
     });
     test('should render unique types', () => {
       const unionModel1 = new ConstrainedStringModel(
@@ -242,7 +242,7 @@ describe('PythonConstrainer', () => {
         constrainedModel: model,
         ...defaultOptions
       });
-      expect(type).toEqual('str');
+      expect(type).toEqual('Union[str]');
     });
   });
 

test/generators/python/__snapshots__/PythonGenerator.spec.ts.snap

@@ -12,7 +12,7 @@ exports[`PythonGenerator Class should handle self reference models 1`] = `
     if 'map_model' in input:
       self._map_model: dict[str, Address] = input['map_model']
     if 'union_model' in input:
-      self._union_model: Address | str = input['union_model']
+      self._union_model: Union[Address, str] = input['union_model']
 
   @property
   def self_model(self) -> Address:
@@ -43,10 +43,10 @@ exports[`PythonGenerator Class should handle self reference models 1`] = `
     self._map_model = map_model
 
   @property
-  def union_model(self) -> Address | str:
+  def union_model(self) -> Union[Address, str]:
     return self._union_model
   @union_model.setter
-  def union_model(self, union_model: Address | str):
+  def union_model(self, union_model: Union[Address, str]):
     self._union_model = union_model
 "
 `;
@@ -76,10 +76,10 @@ exports[`PythonGenerator Class should render \`class\` type 1`] = `
     if 'marriage' in input:
       self._marriage: bool = input['marriage']
     if 'members' in input:
-      self._members: str | float | bool = input['members']
-    self._array_type: List[str | float | Any] = input['array_type']
+      self._members: Union[str, float, bool] = input['members']
+    self._array_type: List[Union[str, float, Any]] = input['array_type']
     if 'additional_properties' in input:
-      self._additional_properties: dict[str, Any | str] = input['additional_properties']
+      self._additional_properties: dict[str, Union[Any, str]] = input['additional_properties']
 
   @property
   def street_name(self) -> str:
@@ -117,24 +117,24 @@ exports[`PythonGenerator Class should render \`class\` type 1`] = `
     self._marriage = marriage
 
   @property
-  def members(self) -> str | float | bool:
+  def members(self) -> Union[str, float, bool]:
     return self._members
   @members.setter
-  def members(self, members: str | float | bool):
+  def members(self, members: Union[str, float, bool]):
     self._members = members
 
   @property
-  def array_type(self) -> List[str | float | Any]:
+  def array_type(self) -> List[Union[str, float, Any]]:
     return self._array_type
   @array_type.setter
-  def array_type(self, array_type: List[str | float | Any]):
+  def array_type(self, array_type: List[Union[str, float, Any]]):
     self._array_type = array_type
 
   @property
-  def additional_properties(self) -> dict[str, Any | str]:
+  def additional_properties(self) -> dict[str, Union[Any, str]]:
     return self._additional_properties
   @additional_properties.setter
-  def additional_properties(self, additional_properties: dict[str, Any | str]):
+  def additional_properties(self, additional_properties: dict[str, Union[Any, str]]):
     self._additional_properties = additional_properties
 "
 `;

test/generators/python/presets/Pydantic.spec.ts

@@ -104,19 +104,44 @@ describe('PYTHON_PYDANTIC_PRESET', () => {
       },
       components: {
         messages: {
+          Garage: {
+            payload: {
+              $ref: '#/components/schemas/Garage'
+            }
+          },
           Vehicle: {
             payload: {
-              oneOf: [
-                { $ref: '#/components/schemas/Car' },
-                { $ref: '#/components/schemas/Truck' }
-              ]
+              $ref: '#/components/schemas/Vehicle'
             }
           }
         },
         schemas: {
+          Garage: {
+            title: 'Garage',
+            type: 'object',
+            properties: {
+              favorite: {
+                $ref: '#/components/schemas/Vehicle'
+              },
+              vehicles: {
+                type: 'array',
+                items: {
+                  $ref: '#/components/schemas/Vehicle'
+                }
+              }
+            }
+          },
           Vehicle: {
             title: 'Vehicle',
             type: 'object',
+            oneOf: [
+              { $ref: '#/components/schemas/Car' },
+              { $ref: '#/components/schemas/Truck' }
+            ]
+          },
+          VehicleBase: {
+            title: 'VehicleBase',
+            type: 'object',
             discriminator: 'vehicleType',
             properties: {
               vehicleType: {
@@ -132,7 +157,7 @@ describe('PYTHON_PYDANTIC_PRESET', () => {
           },
           Car: {
             allOf: [
-              { $ref: '#/components/schemas/Vehicle' },
+              { $ref: '#/components/schemas/VehicleBase' },
               {
                 type: 'object',
                 properties: {
@@ -145,7 +170,7 @@ describe('PYTHON_PYDANTIC_PRESET', () => {
           },
           Truck: {
             allOf: [
-              { $ref: '#/components/schemas/Vehicle' },
+              { $ref: '#/components/schemas/VehicleBase' },
               {
                 type: 'object',
                 properties: {

test/generators/python/presets/__snapshots__/Pydantic.spec.ts.snap

@@ -45,8 +45,46 @@ Array [
 exports[`PYTHON_PYDANTIC_PRESET should render default value for discriminator when using polymorphism 1`] = `
 Array [
   "",
+  "class Garage(BaseModel): 
+  favorite: Optional[Annotated[Union[Car, Truck], Field(discriminator='vehicle_type')]] = Field(default=None)
+  vehicles: Optional[List[Annotated[Union[Car, Truck], Field(discriminator='vehicle_type')]]] = Field(default=None)
+  additional_properties: Optional[dict[str, Any]] = Field(default=None, exclude=True)
+
+  @model_serializer(mode='wrap')
+  def custom_serializer(self, handler):
+    serialized_self = handler(self)
+    additional_properties = getattr(self, \\"additional_properties\\")
+    if additional_properties is not None:
+      for key, value in additional_properties.items():
+        # Never overwrite existing values, to avoid clashes
+        if not hasattr(serialized_self, key):
+          serialized_self[key] = value
+
+    return serialized_self
+
+  @model_validator(mode='before')
+  @classmethod
+  def unwrap_additional_properties(cls, data):
+    if not isinstance(data, dict):
+      data = data.model_dump()
+    json_properties = list(data.keys())
+    known_object_properties = ['favorite', 'vehicles', 'additional_properties']
+    unknown_object_properties = [element for element in json_properties if element not in known_object_properties]
+    # Ignore attempts that validate regular models, only when unknown input is used we add unwrap extensions
+    if len(unknown_object_properties) == 0: 
+      return data
+  
+    known_json_properties = ['favorite', 'vehicles', 'additionalProperties']
+    additional_properties = data.get('additional_properties', {})
+    for obj_key in unknown_object_properties:
+      if not known_json_properties.__contains__(obj_key):
+        additional_properties[obj_key] = data.pop(obj_key, None)
+    data['additional_properties'] = additional_properties
+    return data
+
+",
   "class Car(BaseModel): 
-  vehicle_type: VehicleType = Field(default=VehicleType.CAR, frozen=True, alias='''vehicleType''')
+  vehicle_type: Literal['Car'] = Field(default='Car', frozen=True, alias='''vehicleType''')
   length: Optional[float] = Field(default=None)
   additional_properties: Optional[dict[str, Any]] = Field(default=None, exclude=True)
 
@@ -87,7 +125,7 @@ Array [
   CAR = \\"Car\\"
   TRUCK = \\"Truck\\"",
   "class Truck(BaseModel): 
-  vehicle_type: VehicleType = Field(default=VehicleType.TRUCK, frozen=True, alias='''vehicleType''')
+  vehicle_type: Literal['Truck'] = Field(default='Truck', frozen=True, alias='''vehicleType''')
   length: Optional[float] = Field(default=None)
   additional_properties: Optional[dict[str, Any]] = Field(default=None, exclude=True)