Merge pull request #22 from adrianhdezm/codex/generate-zod-schemas-and-typeddict-for-oneof-object

Handle oneOf during schema generation

Author: adrianhdezm

src/utils/json-schema-to-typed-dict.ts

@@ -92,6 +92,32 @@ export function convertToTypedDict(name: string, schema: OpenAPI.SchemaObject |
       return `${refName}`;
     }
 
+    if ('oneOf' in s && Array.isArray(s.oneOf)) {
+      if (s.oneOf.length === 0) {
+        typingImports.add('Any');
+        return 'Any';
+      }
+      if (s.oneOf.length === 1) {
+        return toType(s.oneOf[0] as any, className);
+      }
+      typingImports.add('Union');
+      const parts = s.oneOf.map((sub, idx) => toType(sub as any, `${className}Option${idx}`)).join(', ');
+      return `Union[${parts}]`;
+    }
+
+    if ('anyOf' in s && Array.isArray(s.anyOf)) {
+      if (s.anyOf.length === 0) {
+        typingImports.add('Any');
+        return 'Any';
+      }
+      if (s.anyOf.length === 1) {
+        return toType(s.anyOf[0] as any, className);
+      }
+      typingImports.add('Union');
+      const parts = s.anyOf.map((sub, idx) => toType(sub as any, `${className}Option${idx}`)).join(', ');
+      return `Union[${parts}]`;
+    }
+
     if (s.enum) {
       typingImports.add('Literal');
       const values = s.enum.map((v) => JSON.stringify(v)).join(', ');

src/utils/json-schema-to-zod.ts

@@ -28,6 +28,30 @@ export function convertSchema(schema: OpenAPI.SchemaObject | OpenAPI.ReferenceOb
       return expr;
     }
 
+    if ('oneOf' in s && Array.isArray(s.oneOf)) {
+      const items = s.oneOf as Array<OpenAPI.SchemaObject | OpenAPI.ReferenceObject>;
+      if (items.length === 0) {
+        return 'z.any()';
+      }
+      if (items.length === 1) {
+        return walk(items[0]!);
+      }
+      const parts = items.map((sub) => walk(sub)).join(', ');
+      return `z.union([${parts}])`;
+    }
+
+    if ('anyOf' in s && Array.isArray(s.anyOf)) {
+      const items = s.anyOf as Array<OpenAPI.SchemaObject | OpenAPI.ReferenceObject>;
+      if (items.length === 0) {
+        return 'z.any()';
+      }
+      if (items.length === 1) {
+        return walk(items[0]!);
+      }
+      const parts = items.map((sub) => walk(sub)).join(', ');
+      return `z.union([${parts}])`;
+    }
+
     if (s.enum) {
       const values = s.enum.map((v) => JSON.stringify(v)).join(', ');
       return `z.enum([${values}])`;

tests/__snapshots__/json-schema-to-typed-dict.spec.ts.snap

@@ -38,3 +38,13 @@ class Place(TypedDict, total=False):
     name: Required[str]
     location: Required[PlaceLocation]"
 `;
+
+exports[`generate-python-dict > handles oneOf with a single ref 1`] = `
+"from typing import TypedDict
+MessageSingle = A"
+`;
+
+exports[`generate-python-dict > handles oneOf with refs 1`] = `
+"from typing import TypedDict, Union
+Message = Union[A, B]"
+`;

tests/json-schema-to-typed-dict.spec.ts

@@ -102,6 +102,53 @@ describe('generate-python-dict', () => {
     expect(content).toMatchSnapshot();
   });
 
+  it('handles oneOf with refs', () => {
+    const doc: OpenAPI.Document = {
+      openapi: '3.1.0',
+      info: { title: 't', version: '1' },
+      paths: {},
+      components: {
+        schemas: {
+          A: { type: 'object', properties: { id: { type: 'string' } } },
+          B: { type: 'object', properties: { value: { type: 'number' } } },
+          Message: { oneOf: [{ $ref: '#/components/schemas/A' }, { $ref: '#/components/schemas/B' }] }
+        }
+      }
+    };
+    const schemas = extractSchemas(doc, null);
+    const { definition, typingImports } = convertToTypedDict('Message', schemas.Message as OpenAPI.SchemaObject);
+    const typingLine = `from typing import ${Array.from(typingImports).join(', ')}`;
+    const content = [typingLine, '', definition, ''].filter(Boolean).join('\n');
+    const tmp = path.join(__dirname, 'tmp_oneof.py');
+    fs.writeFileSync(tmp, content);
+    child_process.execSync(`python3 -m py_compile ${tmp}`);
+    fs.unlinkSync(tmp);
+    expect(content).toMatchSnapshot();
+  });
+
+  it('handles oneOf with a single ref', () => {
+    const doc: OpenAPI.Document = {
+      openapi: '3.1.0',
+      info: { title: 't', version: '1' },
+      paths: {},
+      components: {
+        schemas: {
+          A: { type: 'object', properties: { id: { type: 'string' } } },
+          MessageSingle: { oneOf: [{ $ref: '#/components/schemas/A' }] }
+        }
+      }
+    };
+    const schemas = extractSchemas(doc, null);
+    const { definition, typingImports } = convertToTypedDict('MessageSingle', schemas.MessageSingle as OpenAPI.SchemaObject);
+    const typingLine = `from typing import ${Array.from(typingImports).join(', ')}`;
+    const content = [typingLine, '', definition, ''].filter(Boolean).join('\n');
+    const tmp = path.join(__dirname, 'tmp_oneof_single.py');
+    fs.writeFileSync(tmp, content);
+    child_process.execSync(`python3 -m py_compile ${tmp}`);
+    fs.unlinkSync(tmp);
+    expect(content).toMatchSnapshot();
+  });
+
   it('handles inline object properties', () => {
     const doc: OpenAPI.Document = {
       openapi: '3.1.0',

tests/json-schema-to-zod.spec.ts

@@ -53,6 +53,25 @@ describe('convertSchema', () => {
     expect(imports.has('Base')).toBe(true);
   });
 
+  it('handles oneOf with refs', () => {
+    const schema = {
+      oneOf: [{ $ref: '#/components/schemas/A' }, { $ref: '#/components/schemas/B' }]
+    } as OpenAPI.SchemaObject;
+    const { zodString, imports } = convertSchema(schema);
+    expect(zodString).toBe('z.union([A, B])');
+    expect(imports.has('A')).toBe(true);
+    expect(imports.has('B')).toBe(true);
+  });
+
+  it('handles oneOf with a single ref', () => {
+    const schema = {
+      oneOf: [{ $ref: '#/components/schemas/A' }]
+    } as OpenAPI.SchemaObject;
+    const { zodString, imports } = convertSchema(schema);
+    expect(zodString).toBe('A');
+    expect(imports.has('A')).toBe(true);
+  });
+
   it('converts inline object properties', () => {
     const schema = {
       type: 'object',