feat(python): add Literal type support for const properties in Pydantic preset (#2388)

Author: varundeepsaini

src/generators/python/presets/Pydantic.ts

@@ -5,6 +5,60 @@ import {
 import { PythonOptions } from '../PythonGenerator';
 import { ClassPresetType, PythonPreset } from '../PythonPreset';
 
+function formatPythonConstValue(constValue: unknown): string {
+  if (typeof constValue === 'string') {
+    return `'${constValue}'`;
+  }
+  if (typeof constValue === 'boolean') {
+    return constValue ? 'True' : 'False';
+  }
+  return String(constValue);
+}
+
+function formatLiteralType(constValue: unknown): string {
+  return `Literal[${formatPythonConstValue(constValue)}]`;
+}
+
+function buildFieldArgs(
+  property: ConstrainedObjectPropertyModel,
+  isOptional: boolean,
+  constOptions?: { originalInput: unknown }
+): string[] {
+  const decoratorArgs: string[] = [];
+
+  if (property.property.originalInput['description']) {
+    decoratorArgs.push(
+      `description='''${property.property.originalInput['description']}'''`
+    );
+  }
+
+  if (constOptions) {
+    decoratorArgs.push(
+      `default=${formatPythonConstValue(constOptions.originalInput)}`
+    );
+    decoratorArgs.push('frozen=True');
+  } else if (isOptional) {
+    decoratorArgs.push('default=None');
+  }
+
+  const isUnwrappedDict =
+    property.property instanceof ConstrainedDictionaryModel &&
+    property.property.serializationType === 'unwrap';
+
+  if (isUnwrappedDict) {
+    decoratorArgs.push('exclude=True');
+  }
+
+  if (
+    property.propertyName !== property.unconstrainedPropertyName &&
+    !isUnwrappedDict
+  ) {
+    decoratorArgs.push(`alias='''${property.unconstrainedPropertyName}'''`);
+  }
+
+  return decoratorArgs;
+}
+
 const PYTHON_PYDANTIC_CLASS_PRESET: ClassPresetType<PythonOptions> = {
   async self({ renderer, model }) {
     renderer.dependencyManager.addDependency(
@@ -22,61 +76,25 @@ const PYTHON_PYDANTIC_CLASS_PRESET: ClassPresetType<PythonOptions> = {
     );
   },
   property({ property, model, renderer }) {
-    let type = property.property.type;
     const propertyName = property.propertyName;
-
+    const constOptions = property.property.options.const;
     const isOptional =
       !property.required || property.property.options.isNullable === true;
+
+    let type = property.property.type;
     if (isOptional) {
       type = `Optional[${type}]`;
     }
-    if (
-      property.property.options.const &&
-      model.options.discriminator?.discriminator ===
-        property.unconstrainedPropertyName
-    ) {
+    if (constOptions) {
       renderer.dependencyManager.addDependency('from typing import Literal');
-      type = `Literal['${property.property.options.const.originalInput}']`;
+      type = formatLiteralType(constOptions.originalInput);
     }
     type = renderer.renderPropertyType({
       modelType: model.type,
       propertyType: type
     });
 
-    const decoratorArgs: string[] = [];
-
-    if (property.property.originalInput['description']) {
-      decoratorArgs.push(
-        `description='''${property.property.originalInput['description']}'''`
-      );
-    }
-    if (isOptional) {
-      decoratorArgs.push('default=None');
-    }
-    if (property.property.options.const) {
-      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 (
-      property.property instanceof ConstrainedDictionaryModel &&
-      property.property.serializationType === 'unwrap'
-    ) {
-      decoratorArgs.push('exclude=True');
-    }
-    if (
-      property.propertyName !== property.unconstrainedPropertyName &&
-      (!(property.property instanceof ConstrainedDictionaryModel) ||
-        property.property.serializationType !== 'unwrap')
-    ) {
-      decoratorArgs.push(`alias='''${property.unconstrainedPropertyName}'''`);
-    }
+    const decoratorArgs = buildFieldArgs(property, isOptional, constOptions);
 
     return `${propertyName}: ${type} = Field(${decoratorArgs.join(', ')})`;
   },

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

@@ -204,4 +204,25 @@ describe('PYTHON_PYDANTIC_PRESET', () => {
     const models = await generator.generate(doc);
     expect(models.map((model) => model.result)).toMatchSnapshot();
   });
+
+  test('should render Literal type for const properties', async () => {
+    const doc = {
+      title: 'ConstTest',
+      type: 'object',
+      properties: {
+        country: {
+          const: 'United States of America'
+        },
+        version: {
+          const: 42
+        },
+        isActive: {
+          const: true
+        }
+      }
+    };
+
+    const models = await generator.generate(doc);
+    expect(models.map((model) => model.result)).toMatchSnapshot();
+  });
 });

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

@@ -42,6 +42,50 @@ Array [
 ]
 `;
 
+exports[`PYTHON_PYDANTIC_PRESET should render Literal type for const properties 1`] = `
+Array [
+  "class ConstTest(BaseModel): 
+  country: Literal['United States of America'] = Field(default='United States of America', frozen=True)
+  version: Literal[42] = Field(default=42, frozen=True)
+  is_active: Literal[True] = Field(default=True, frozen=True, alias='''isActive''')
+  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 key in serialized_self:
+          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 = ['country', 'version', 'is_active', '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 = ['country', 'version', 'isActive', '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
+
+",
+]
+`;
+
 exports[`PYTHON_PYDANTIC_PRESET should render default value for discriminator when using polymorphism 1`] = `
 Array [
   "",