From 7cf76c04c85f0b5353de321d7cd865afb0261479 Mon Sep 17 00:00:00 2001 From: GrantRheingold Date: Wed, 18 Jun 2025 13:55:45 -0400 Subject: [PATCH 1/6] union implementation --- .../pydantic-model/src/v2/UnionGenerator.ts | 88 +++++++++++++++++++ .../pydantic-model/src/v2/generateV2Models.ts | 5 +- 2 files changed, 92 insertions(+), 1 deletion(-) create mode 100644 generators/python-v2/pydantic-model/src/v2/UnionGenerator.ts diff --git a/generators/python-v2/pydantic-model/src/v2/UnionGenerator.ts b/generators/python-v2/pydantic-model/src/v2/UnionGenerator.ts new file mode 100644 index 00000000000..c9834662086 --- /dev/null +++ b/generators/python-v2/pydantic-model/src/v2/UnionGenerator.ts @@ -0,0 +1,88 @@ +import { RelativeFilePath } from "@fern-api/fs-utils"; +import { python } from "@fern-api/python-ast"; +import { WriteablePythonFile } from "@fern-api/python-base"; + +import { SingleUnionType, TypeDeclaration, TypeId, UnionTypeDeclaration } from "@fern-fern/ir-sdk/api"; + +import { PydanticModelGeneratorContext } from "../ModelGeneratorContext"; + +export class UnionGenerator { + constructor( + private readonly typeId: TypeId, + private readonly context: PydanticModelGeneratorContext, + private readonly typeDeclaration: TypeDeclaration, + private readonly unionDeclaration: UnionTypeDeclaration + ) {} + + public doGenerate(): WriteablePythonFile { + const path = this.context.getModulePathForId(this.typeId); + const filename = this.context.getSnakeCaseSafeName(this.typeDeclaration.name.name); + const file = python.file({ path }); + + // Add imports + file.addStatement(python.starImport({ modulePath: "typing", names: ["Union"] })); + file.addStatement(python.starImport({ modulePath: "typing_extensions", names: ["Literal"] })); + + // Generate variant classes + this.unionDeclaration.types.forEach((variant: SingleUnionType) => { + const variantClass = python.class_({ + name: `${this.context.getPascalCaseSafeName(this.typeDeclaration.name.name)}_${this.context.getPascalCaseSafeName(variant.name.name)}`, + extends_: [ + ...this.unionDeclaration.extends.map((type) => + this.context.pythonTypeMapper.convertToClassReference(type) + ), + this.context.pythonTypeMapper.convertToClassReference(variant.type) + ], + fields: [ + python.field({ + name: this.unionDeclaration.discriminant.name, + type: python.typeHint.literal(variant.name.wireValue) + }), + ...this.unionDeclaration.baseProperties.map((prop) => + python.field({ + name: prop.name.name, + type: this.context.pythonTypeMapper.convertToTypeHint(prop.valueType) + }) + ) + ] + }); + + // Add Config class + variantClass.add( + python.class_({ + name: "Config", + fields: [ + python.field({ + name: "allow_population_by_field_name", + initializer: python.TypeInstantiation.bool(true) + }) + ] + }) + ); + + file.addStatement(variantClass); + }); + + // Add union type + file.addStatement( + python.codeBlock((writer) => { + writer.write(`${this.context.getPascalCaseSafeName(this.typeDeclaration.name.name)} = Union[`); + this.unionDeclaration.types.forEach((variant, index) => { + if (index > 0) { + writer.write(", "); + } + writer.write( + `${this.context.getPascalCaseSafeName(this.typeDeclaration.name.name)}_${this.context.getPascalCaseSafeName(variant.name.name)}` + ); + }); + writer.write("]"); + }) + ); + + return new WriteablePythonFile({ + contents: file, + directory: RelativeFilePath.of(path.join("/")), + filename + }); + } +} diff --git a/generators/python-v2/pydantic-model/src/v2/generateV2Models.ts b/generators/python-v2/pydantic-model/src/v2/generateV2Models.ts index 109ed42d0b4..a81ad02829f 100644 --- a/generators/python-v2/pydantic-model/src/v2/generateV2Models.ts +++ b/generators/python-v2/pydantic-model/src/v2/generateV2Models.ts @@ -2,6 +2,7 @@ import { WriteablePythonFile } from "@fern-api/python-base"; import { PydanticModelGeneratorContext } from "../ModelGeneratorContext"; import { ObjectGenerator } from "./ObjectGenerator"; +import { UnionGenerator } from "./UnionGenerator"; import { WrappedAliasGenerator } from "./WrappedAliasGenerator"; export function generateV2Models({ context }: { context: PydanticModelGeneratorContext }): WriteablePythonFile[] { @@ -16,7 +17,9 @@ export function generateV2Models({ context }: { context: PydanticModelGeneratorC return new ObjectGenerator(typeId, context, typeDeclaration, objectTypDeclaration).doGenerate(); }, undiscriminatedUnion: () => undefined, - union: () => undefined, + union: (unionTypeDeclaration) => { + return new UnionGenerator(typeId, context, typeDeclaration, unionTypeDeclaration).doGenerate(); + }, _other: () => undefined }); if (file != null) { From 65a4274c6fbb7ee846726e04a5d0158a0eb170de Mon Sep 17 00:00:00 2001 From: GrantRheingold Date: Wed, 18 Jun 2025 14:01:53 -0400 Subject: [PATCH 2/6] progress on union --- .../pydantic-model/src/v2/UnionGenerator.ts | 32 ++++--------------- .../types/resources/union/types/animal.py | 14 ++++++++ 2 files changed, 20 insertions(+), 26 deletions(-) create mode 100644 seed/pydantic-v2/exhaustive/resources/types/resources/union/types/animal.py diff --git a/generators/python-v2/pydantic-model/src/v2/UnionGenerator.ts b/generators/python-v2/pydantic-model/src/v2/UnionGenerator.ts index c9834662086..648347b03a2 100644 --- a/generators/python-v2/pydantic-model/src/v2/UnionGenerator.ts +++ b/generators/python-v2/pydantic-model/src/v2/UnionGenerator.ts @@ -20,43 +20,25 @@ export class UnionGenerator { const file = python.file({ path }); // Add imports - file.addStatement(python.starImport({ modulePath: "typing", names: ["Union"] })); - file.addStatement(python.starImport({ modulePath: "typing_extensions", names: ["Literal"] })); + file.addReference(python.reference({ name: "Union", modulePath: ["typing"] })); + file.addReference(python.reference({ name: "Literal", modulePath: ["typing_extensions"] })); // Generate variant classes this.unionDeclaration.types.forEach((variant: SingleUnionType) => { const variantClass = python.class_({ - name: `${this.context.getPascalCaseSafeName(this.typeDeclaration.name.name)}_${this.context.getPascalCaseSafeName(variant.name.name)}`, + name: "name", extends_: [ ...this.unionDeclaration.extends.map((type) => this.context.pythonTypeMapper.convertToClassReference(type) - ), - this.context.pythonTypeMapper.convertToClassReference(variant.type) - ], - fields: [ - python.field({ - name: this.unionDeclaration.discriminant.name, - type: python.typeHint.literal(variant.name.wireValue) - }), - ...this.unionDeclaration.baseProperties.map((prop) => - python.field({ - name: prop.name.name, - type: this.context.pythonTypeMapper.convertToTypeHint(prop.valueType) - }) ) + //this.context.pythonTypeMapper.convertToClassReference(variant.type) ] }); // Add Config class variantClass.add( python.class_({ - name: "Config", - fields: [ - python.field({ - name: "allow_population_by_field_name", - initializer: python.TypeInstantiation.bool(true) - }) - ] + name: "Config" }) ); @@ -71,9 +53,7 @@ export class UnionGenerator { if (index > 0) { writer.write(", "); } - writer.write( - `${this.context.getPascalCaseSafeName(this.typeDeclaration.name.name)}_${this.context.getPascalCaseSafeName(variant.name.name)}` - ); + writer.write("name"); }); writer.write("]"); }) diff --git a/seed/pydantic-v2/exhaustive/resources/types/resources/union/types/animal.py b/seed/pydantic-v2/exhaustive/resources/types/resources/union/types/animal.py new file mode 100644 index 00000000000..fba77a87dc3 --- /dev/null +++ b/seed/pydantic-v2/exhaustive/resources/types/resources/union/types/animal.py @@ -0,0 +1,14 @@ +from typing import Union +from typing_extensions import Literal + +class name: + class Config: + pass + + +class name: + class Config: + pass + + +Animal = Union[name, name] From 8c1ce653a2d1d694a3634ecdd68047694c89e38b Mon Sep 17 00:00:00 2001 From: GrantRheingold Date: Wed, 18 Jun 2025 15:17:03 -0400 Subject: [PATCH 3/6] updated union logic --- .../pydantic-model/src/v2/UnionGenerator.ts | 43 +++++++++++++++---- .../types/resources/union/types/animal.py | 17 +++++--- 2 files changed, 45 insertions(+), 15 deletions(-) diff --git a/generators/python-v2/pydantic-model/src/v2/UnionGenerator.ts b/generators/python-v2/pydantic-model/src/v2/UnionGenerator.ts index 648347b03a2..19278cf3c6d 100644 --- a/generators/python-v2/pydantic-model/src/v2/UnionGenerator.ts +++ b/generators/python-v2/pydantic-model/src/v2/UnionGenerator.ts @@ -2,7 +2,7 @@ import { RelativeFilePath } from "@fern-api/fs-utils"; import { python } from "@fern-api/python-ast"; import { WriteablePythonFile } from "@fern-api/python-base"; -import { SingleUnionType, TypeDeclaration, TypeId, UnionTypeDeclaration } from "@fern-fern/ir-sdk/api"; +import { Literal, SingleUnionType, TypeDeclaration, TypeId, UnionTypeDeclaration } from "@fern-fern/ir-sdk/api"; import { PydanticModelGeneratorContext } from "../ModelGeneratorContext"; @@ -21,27 +21,37 @@ export class UnionGenerator { // Add imports file.addReference(python.reference({ name: "Union", modulePath: ["typing"] })); - file.addReference(python.reference({ name: "Literal", modulePath: ["typing_extensions"] })); + //file.addReference(python.reference({ name: "Literal", modulePath: ["typing_extensions"] })); // Generate variant classes this.unionDeclaration.types.forEach((variant: SingleUnionType) => { + file.addReference( + python.reference({ + name: `${this.context.getPascalCaseSafeName(variant.discriminantValue.name)}`, + modulePath: [`.${this.context.getPascalCaseSafeName(variant.discriminantValue.name).toLowerCase()}`] + }) + ); const variantClass = python.class_({ - name: "name", + name: `${this.context.getPascalCaseSafeName(this.typeDeclaration.name.name)}_${this.context.getPascalCaseSafeName(variant.discriminantValue.name)}(${this.context.getPascalCaseSafeName(variant.discriminantValue.name)})`, extends_: [ ...this.unionDeclaration.extends.map((type) => this.context.pythonTypeMapper.convertToClassReference(type) ) - //this.context.pythonTypeMapper.convertToClassReference(variant.type) ] }); - // Add Config class + // Add the literal type field variantClass.add( - python.class_({ - name: "Config" + python.field({ + name: this.context.getSnakeCaseSafeName(this.typeDeclaration.name.name), + type: python.Type.literal( + `${this.context.getPascalCaseSafeName(variant.discriminantValue.name).toLowerCase()}` + ) }) ); + variantClass.add(this.getConfigClass()); + file.addStatement(variantClass); }); @@ -53,7 +63,9 @@ export class UnionGenerator { if (index > 0) { writer.write(", "); } - writer.write("name"); + writer.write( + `${this.context.getPascalCaseSafeName(this.typeDeclaration.name.name)}_${this.context.getPascalCaseSafeName(variant.discriminantValue.name)}` + ); }); writer.write("]"); }) @@ -65,4 +77,19 @@ export class UnionGenerator { filename }); } + + private getConfigClass(): python.Class { + const configClass = python.class_({ + name: "Config" + }); + + configClass.addField( + python.field({ + name: "allow_population_by_field_name", + initializer: python.TypeInstantiation.bool(true) + }) + ); + + return configClass; + } } diff --git a/seed/pydantic-v2/exhaustive/resources/types/resources/union/types/animal.py b/seed/pydantic-v2/exhaustive/resources/types/resources/union/types/animal.py index fba77a87dc3..5bc5b0a3b2b 100644 --- a/seed/pydantic-v2/exhaustive/resources/types/resources/union/types/animal.py +++ b/seed/pydantic-v2/exhaustive/resources/types/resources/union/types/animal.py @@ -1,14 +1,17 @@ -from typing import Union -from typing_extensions import Literal +from typing import Union, Literal +from .dog import Dog +from .cat import Cat -class name: +class Animal_Dog(Dog): + animal: Literal["dog"] class Config: - pass + allow_population_by_field_name = True -class name: +class Animal_Cat(Cat): + animal: Literal["cat"] class Config: - pass + allow_population_by_field_name = True -Animal = Union[name, name] +Animal = Union[Animal_Dog, Animal_Cat] From de3328b10317e08efa5290f0f23963c6c6388e1a Mon Sep 17 00:00:00 2001 From: GrantRheingold Date: Wed, 18 Jun 2025 15:45:44 -0400 Subject: [PATCH 4/6] update union --- generators/python-v2/pydantic-model/src/v2/UnionGenerator.ts | 1 - .../exhaustive/resources/types/resources/union/types/animal.py | 1 + 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/generators/python-v2/pydantic-model/src/v2/UnionGenerator.ts b/generators/python-v2/pydantic-model/src/v2/UnionGenerator.ts index 19278cf3c6d..83178a9769b 100644 --- a/generators/python-v2/pydantic-model/src/v2/UnionGenerator.ts +++ b/generators/python-v2/pydantic-model/src/v2/UnionGenerator.ts @@ -21,7 +21,6 @@ export class UnionGenerator { // Add imports file.addReference(python.reference({ name: "Union", modulePath: ["typing"] })); - //file.addReference(python.reference({ name: "Literal", modulePath: ["typing_extensions"] })); // Generate variant classes this.unionDeclaration.types.forEach((variant: SingleUnionType) => { diff --git a/seed/pydantic-v2/exhaustive/resources/types/resources/union/types/animal.py b/seed/pydantic-v2/exhaustive/resources/types/resources/union/types/animal.py index 5bc5b0a3b2b..1d32763db52 100644 --- a/seed/pydantic-v2/exhaustive/resources/types/resources/union/types/animal.py +++ b/seed/pydantic-v2/exhaustive/resources/types/resources/union/types/animal.py @@ -1,4 +1,5 @@ from typing import Union, Literal +from typing_extensions import Literal as TypingExtensionsLiteral from .dog import Dog from .cat import Cat From 4c08e419c60a3debbdd649cec3f6ba30ad37cc49 Mon Sep 17 00:00:00 2001 From: GrantRheingold Date: Wed, 18 Jun 2025 20:51:35 -0400 Subject: [PATCH 5/6] update imports --- generators/python-v2/pydantic-model/src/v2/UnionGenerator.ts | 1 + 1 file changed, 1 insertion(+) diff --git a/generators/python-v2/pydantic-model/src/v2/UnionGenerator.ts b/generators/python-v2/pydantic-model/src/v2/UnionGenerator.ts index 83178a9769b..3d95c702a21 100644 --- a/generators/python-v2/pydantic-model/src/v2/UnionGenerator.ts +++ b/generators/python-v2/pydantic-model/src/v2/UnionGenerator.ts @@ -21,6 +21,7 @@ export class UnionGenerator { // Add imports file.addReference(python.reference({ name: "Union", modulePath: ["typing"] })); + file.addReference(python.reference({ name: "Literal", modulePath: ["typing"] })); // Generate variant classes this.unionDeclaration.types.forEach((variant: SingleUnionType) => { From 1ba3debdb33dbef541bd30b0369c0ab73853e789 Mon Sep 17 00:00:00 2001 From: GrantRheingold Date: Wed, 18 Jun 2025 21:15:59 -0400 Subject: [PATCH 6/6] remove unused import --- generators/python-v2/pydantic-model/src/v2/UnionGenerator.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/generators/python-v2/pydantic-model/src/v2/UnionGenerator.ts b/generators/python-v2/pydantic-model/src/v2/UnionGenerator.ts index 3d95c702a21..42f1c594c75 100644 --- a/generators/python-v2/pydantic-model/src/v2/UnionGenerator.ts +++ b/generators/python-v2/pydantic-model/src/v2/UnionGenerator.ts @@ -2,7 +2,7 @@ import { RelativeFilePath } from "@fern-api/fs-utils"; import { python } from "@fern-api/python-ast"; import { WriteablePythonFile } from "@fern-api/python-base"; -import { Literal, SingleUnionType, TypeDeclaration, TypeId, UnionTypeDeclaration } from "@fern-fern/ir-sdk/api"; +import { SingleUnionType, TypeDeclaration, TypeId, UnionTypeDeclaration } from "@fern-fern/ir-sdk/api"; import { PydanticModelGeneratorContext } from "../ModelGeneratorContext";