Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
102 changes: 60 additions & 42 deletions src/generators/python/presets/Pydantic.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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(', ')})`;
},
Expand Down
21 changes: 21 additions & 0 deletions test/generators/python/presets/Pydantic.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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();
});
});
Original file line number Diff line number Diff line change
Expand Up @@ -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 [
"",
Expand Down