feat!: render python union in pydantic in the pre 3.10 way (#1626)

asyncapi · Nov 21, 2023 · 7ac3259 · 7ac3259
1 parent 41c2daf
commit 7ac3259
Show file tree

Hide file tree

Showing 4 changed files with 103 additions and 5 deletions.
diff --git a/docs/migrations/version-2-to-3.md b/docs/migrations/version-2-to-3.md
@@ -233,7 +233,43 @@ Is not affected by this change.
 
 ### Python
 
-Is not affected by this change.
+#### Union type for the Pydantic preset supports Python pre 3.10
+
+Modelina used to use the newer way of representing unions in Python by using the `|` operator. In the Pydantic preset, this is now adjusted to support Python pre 3.10 by using `Union[Model1, Model2]` instead:
+
+```yaml
+title: UnionTest
+type: object
+  properties:
+    unionTest:
+      oneOf:
+        - title: Union1
+          type: object
+          properties:
+            testProp1:
+              type: string
+        - title: Union2
+          type: object
+          properties:
+            testProp2:
+              type: string
+```
+
+will generate
+
+```python
+class UnionTest(BaseModel):
+  unionTest: Optional[Union[Union1, Union2]] = Field()
+  additionalProperties: Optional[dict[Any, Any]] = Field()
+
+class Union1(BaseModel):
+  testProp1: Optional[str] = Field()
+  additionalProperties: Optional[dict[Any, Any]] = Field()
+
+class Union2(BaseModel):
+  testProp2: Optional[str] = Field()
+  additionalProperties: Optional[dict[Any, Any]] = Field()
+```
 
 ### Go
 

diff --git a/src/generators/python/presets/Pydantic.ts b/src/generators/python/presets/Pydantic.ts
@@ -1,10 +1,11 @@
+import { ConstrainedUnionModel } 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'
+      'from typing import Optional, Any, Union'
     );
     renderer.dependencyManager.addDependency(
       'from pydantic import BaseModel, Field'
@@ -18,9 +19,19 @@ const PYTHON_PYDANTIC_CLASS_PRESET: ClassPresetType<PythonOptions> = {
     );
   },
   property(params) {
-    const type = params.property.required
-      ? params.property.property.type
-      : `Optional[${params.property.property.type}]`;
+    let type = params.property.property.type;
+
+    if (params.property.property instanceof ConstrainedUnionModel) {
+      const unionTypes = params.property.property.union.map(
+        (unionModel) => unionModel.type
+      );
+      type = `Union[${unionTypes.join(', ')}]`;
+    }
+
+    if (!params.property.required) {
+      type = `Optional[${type}]`;
+    }
+
     const alias = params.property.property.originalInput['description']
       ? `alias='''${params.property.property.originalInput['description']}'''`
       : '';

diff --git a/test/generators/python/presets/Pydantic.spec.ts b/test/generators/python/presets/Pydantic.spec.ts
@@ -31,4 +31,38 @@ describe('PYTHON_PYDANTIC_PRESET', () => {
     expect(models).toHaveLength(1);
     expect(models[0].result).toMatchSnapshot();
   });
+
+  test('should render union to support Python < 3.10', async () => {
+    const doc = {
+      title: 'UnionTest',
+      type: 'object',
+      properties: {
+        unionTest: {
+          oneOf: [
+            {
+              title: 'Union1',
+              type: 'object',
+              properties: {
+                testProp1: {
+                  type: 'string'
+                }
+              }
+            },
+            {
+              title: 'Union2',
+              type: 'object',
+              properties: {
+                testProp2: {
+                  type: 'string'
+                }
+              }
+            }
+          ]
+        }
+      }
+    };
+
+    const models = await generator.generate(doc);
+    expect(models.map((model) => model.result)).toMatchSnapshot();
+  });
 });
diff --git a/test/generators/python/presets/__snapshots__/Pydantic.spec.ts.snap b/test/generators/python/presets/__snapshots__/Pydantic.spec.ts.snap
@@ -9,3 +9,20 @@ exports[`PYTHON_PYDANTIC_PRESET should render pydantic for class 1`] = `
   additionalProperties: Optional[dict[Any, Any]] = Field()
 "
 `;
+
+exports[`PYTHON_PYDANTIC_PRESET should render union to support Python < 3.10 1`] = `
+Array [
+  "class UnionTest(BaseModel): 
+  unionTest: Optional[Union[Union1, Union2]] = Field()
+  additionalProperties: Optional[dict[Any, Any]] = Field()
+",
+  "class Union1(BaseModel): 
+  testProp1: Optional[str] = Field()
+  additionalProperties: Optional[dict[Any, Any]] = Field()
+",
+  "class Union2(BaseModel): 
+  testProp2: Optional[str] = Field()
+  additionalProperties: Optional[dict[Any, Any]] = Field()
+",
+]
+`;