Merge pull request #19 from adrianhdezm/codex/add-tests-for-json-schema-allof-functionality

feat(utils): handle allOf in schema conversions

Author: adrianhdezm

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

@@ -8,6 +8,29 @@ export interface TypedDictResult {
 export function convertToTypedDict(name: string, schema: OpenAPI.SchemaObject | OpenAPI.ReferenceObject): TypedDictResult {
   const typingImports = new Set<string>(['TypedDict']);
 
+  function flatten(s: OpenAPI.SchemaObject | OpenAPI.ReferenceObject): { bases: string[]; schema: OpenAPI.SchemaObject | null } {
+    if ('$ref' in s) {
+      const match = s.$ref.match(/^#\/components\/schemas\/(.+)$/);
+      return { bases: [match?.[1] ?? s.$ref], schema: null };
+    }
+
+    if ('allOf' in s && Array.isArray(s.allOf)) {
+      const bases: string[] = [];
+      const merged: OpenAPI.SchemaObject = { type: 'object', properties: {}, required: [] };
+      for (const sub of s.allOf as Array<OpenAPI.SchemaObject | OpenAPI.ReferenceObject>) {
+        const res = flatten(sub);
+        bases.push(...res.bases);
+        if (res.schema) {
+          Object.assign(merged.properties!, res.schema.properties);
+          merged.required = Array.from(new Set([...(merged.required ?? []), ...(res.schema.required ?? [])]));
+        }
+      }
+      return { bases, schema: merged };
+    }
+
+    return { bases: [], schema: s };
+  }
+
   function toType(s: OpenAPI.SchemaObject | OpenAPI.ReferenceObject): string {
     if ('$ref' in s) {
       const match = s.$ref.match(/^#\/components\/schemas\/(.+)$/);
@@ -43,10 +66,12 @@ export function convertToTypedDict(name: string, schema: OpenAPI.SchemaObject |
     }
   }
 
+  const { bases, schema: flat } = flatten(schema);
+
   let definition = '';
-  if (!('$ref' in schema) && schema.type === 'object') {
-    const props = schema.properties ?? {};
-    const required = new Set(schema.required ?? []);
+  if (flat && flat.type === 'object') {
+    const props = flat.properties ?? {};
+    const required = new Set(flat.required ?? []);
     const fields: string[] = [];
     const attrLines: string[] = [];
     for (const [key, value] of Object.entries(props)) {
@@ -65,13 +90,14 @@ export function convertToTypedDict(name: string, schema: OpenAPI.SchemaObject |
     if (fields.length === 0) {
       fields.push('    pass');
     }
-    const header = [`class ${name}(TypedDict, total=False):`];
+    const baseList = bases.length > 0 ? `${bases.join(', ')}, ` : '';
+    const header = [`class ${name}(${baseList}TypedDict, total=False):`];
     const docLines: string[] = [];
-    if (schema.description) {
-      docLines.push((schema.description as string).replace(/\n/g, ' '));
+    if (flat.description) {
+      docLines.push((flat.description as string).replace(/\n/g, ' '));
     }
     if (attrLines.length > 0) {
-      if (schema.description) docLines.push('');
+      if (flat.description) docLines.push('');
       docLines.push('Attributes:');
       for (const line of attrLines) {
         docLines.push(`    ${line}`);

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

@@ -16,6 +16,18 @@ export function convertSchema(schema: OpenAPI.SchemaObject | OpenAPI.ReferenceOb
       return name;
     }
 
+    if ('allOf' in s && Array.isArray(s.allOf)) {
+      const items = s.allOf as Array<OpenAPI.SchemaObject | OpenAPI.ReferenceObject>;
+      if (items.length === 0) {
+        return 'z.any()';
+      }
+      let expr = walk(items[0]!);
+      for (const sub of items.slice(1)) {
+        expr = `z.intersection(${expr}, ${walk(sub)})`;
+      }
+      return expr;
+    }
+
     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

@@ -21,3 +21,9 @@ exports[`generate-python-dict > generates enums and nested arrays 1`] = `
 "from typing import TypedDict, List
 Wrapper = List[Status]"
 `;
+
+exports[`generate-python-dict > handles allOf with refs and properties 1`] = `
+"from typing import TypedDict, Required
+class Derived(Base, TypedDict, total=False):
+    extra: Required[str]"
+`;

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

@@ -74,6 +74,34 @@ describe('generate-python-dict', () => {
     expect(definition).toMatchSnapshot();
   });
 
+  it('handles allOf with refs and properties', () => {
+    const doc: OpenAPI.Document = {
+      openapi: '3.1.0',
+      info: { title: 't', version: '1' },
+      paths: {},
+      components: {
+        schemas: {
+          Base: { type: 'object', properties: { id: { type: 'string' } }, required: ['id'] },
+          Derived: {
+            allOf: [
+              { $ref: '#/components/schemas/Base' },
+              { type: 'object', properties: { extra: { type: 'string' } }, required: ['extra'] }
+            ]
+          }
+        }
+      }
+    };
+    const schemas = extractSchemas(doc, null);
+    const { definition, typingImports } = convertToTypedDict('Derived', schemas.Derived 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_allof.py');
+    fs.writeFileSync(tmp, content);
+    child_process.execSync(`python3 -m py_compile ${tmp}`);
+    fs.unlinkSync(tmp);
+    expect(content).toMatchSnapshot();
+  });
+
   it('filters schemas by path prefixes', () => {
     const doc: OpenAPI.Document = {
       openapi: '3.1.0',

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

@@ -43,4 +43,13 @@ describe('convertSchema', () => {
     const { zodString } = convertSchema(schema);
     expect(zodString).toBe('z.object({\n  id: z.string().meta({ description: "identifier" })\n})');
   });
+
+  it('handles allOf with refs and properties', () => {
+    const schema = {
+      allOf: [{ $ref: '#/components/schemas/Base' }, { type: 'object', properties: { extra: { type: 'string' } }, required: ['extra'] }]
+    } as OpenAPI.SchemaObject;
+    const { zodString, imports } = convertSchema(schema);
+    expect(zodString).toBe('z.intersection(Base, z.object({\n  extra: z.string()\n}))');
+    expect(imports.has('Base')).toBe(true);
+  });
 });