Merge pull request #237 from alisabzevari/zod-unions

Code-Hex · web-flow · commit efe86281d2f7 · 2022-11-23T12:36:13.000+09:00
Support union types for zod generator
diff --git a/src/myzod/index.ts b/src/myzod/index.ts
@@ -9,6 +9,7 @@ import {
   ObjectTypeDefinitionNode,
   EnumTypeDefinitionNode,
   FieldDefinitionNode,
+  UnionTypeDefinitionNode,
 } from 'graphql';
 import { DeclarationBlock, indent } from '@graphql-codegen/visitor-plugin-common';
 import { TsVisitor } from '@graphql-codegen/typescript';
@@ -89,6 +90,22 @@ export const MyZodSchemaVisitor = (schema: GraphQLSchema, config: ValidationSche
         .withName(`${enumname}Schema`)
         .withContent(`myzod.enum(${enumname})`).string;
     },
+    UnionTypeDefinition: (node: UnionTypeDefinitionNode) => {
+      const unionName = tsVisitor.convertName(node.name.value);
+      const unionElements = node.types?.map(t => `${tsVisitor.convertName(t.name.value)}Schema()`).join(', ');
+      const unionElementsCount = node.types?.length ?? 0;
+
+      const union =
+        unionElementsCount > 1 ? indent(`return myzod.union([${unionElements}])`) : indent(`return ${unionElements}`);
+
+      const result = new DeclarationBlock({})
+        .export()
+        .asKind('function')
+        .withName(`${unionName}Schema()`)
+        .withBlock(union);
+
+      return result.string;
+    },
   };
 };
 
@@ -181,6 +198,11 @@ const generateNameNodeMyZodSchema = (
     return `${enumName}Schema`;
   }
 
+  if (typ?.astNode?.kind === 'UnionTypeDefinition') {
+    const enumName = tsVisitor.convertName(typ.astNode.name.value);
+    return `${enumName}Schema()`;
+  }
+
   return myzod4Scalar(config, tsVisitor, node.value);
 };
 
diff --git a/src/yup/index.ts b/src/yup/index.ts
@@ -9,6 +9,7 @@ import {
   EnumTypeDefinitionNode,
   ObjectTypeDefinitionNode,
   FieldDefinitionNode,
+  UnionTypeDefinitionNode,
 } from 'graphql';
 import { DeclarationBlock, indent } from '@graphql-codegen/visitor-plugin-common';
 import { TsVisitor } from '@graphql-codegen/typescript';
@@ -28,7 +29,17 @@ export const YupSchemaVisitor = (schema: GraphQLSchema, config: ValidationSchema
       }
       return [importYup];
     },
-    initialEmit: (): string => '',
+    initialEmit: (): string =>
+      new DeclarationBlock({})
+        .asKind('function')
+        .withName('union<T>(...schemas: ReadonlyArray<yup.SchemaOf<T>>): yup.BaseSchema<T>')
+        .withBlock(
+          [
+            indent('return yup.mixed().test({'),
+            indent('test: (value) => schemas.some((schema) => schema.isValidSync(value))', 2),
+            indent('})'),
+          ].join('\n')
+        ).string,
     InputObjectTypeDefinition: (node: InputObjectTypeDefinitionNode) => {
       const name = tsVisitor.convertName(node.name.value);
       importTypes.push(name);
@@ -89,6 +100,24 @@ export const YupSchemaVisitor = (schema: GraphQLSchema, config: ValidationSchema
         .withName(`${enumname}Schema`)
         .withContent(`yup.mixed().oneOf([${values}])`).string;
     },
+    UnionTypeDefinition: (node: UnionTypeDefinitionNode) => {
+      const unionName = tsVisitor.convertName(node.name.value);
+      const unionElements = node.types?.map(t => `${tsVisitor.convertName(t.name.value)}Schema()`).join(', ');
+      const unionElementsCount = node.types?.length ?? 0;
+
+      const union =
+        unionElementsCount > 1
+          ? indent(`return union<${unionName}>(${unionElements})`)
+          : indent(`return ${unionElements}`);
+
+      const result = new DeclarationBlock({})
+        .export()
+        .asKind('function')
+        .withName(`${unionName}Schema(): yup.BaseSchema<${unionName}>`)
+        .withBlock(union);
+
+      return result.string;
+    },
     // ScalarTypeDefinition: (node) => {
     //   const decl = new DeclarationBlock({})
     //     .export()
diff --git a/src/zod/index.ts b/src/zod/index.ts
@@ -8,6 +8,7 @@ import {
   InputObjectTypeDefinitionNode,
   ObjectTypeDefinitionNode,
   EnumTypeDefinitionNode,
+  UnionTypeDefinitionNode,
   FieldDefinitionNode,
 } from 'graphql';
 import { DeclarationBlock, indent } from '@graphql-codegen/visitor-plugin-common';
@@ -100,6 +101,22 @@ export const ZodSchemaVisitor = (schema: GraphQLSchema, config: ValidationSchema
         .withName(`${enumname}Schema`)
         .withContent(`z.nativeEnum(${enumname})`).string;
     },
+    UnionTypeDefinition: (node: UnionTypeDefinitionNode) => {
+      const unionName = tsVisitor.convertName(node.name.value);
+      const unionElements = node.types?.map(t => `${tsVisitor.convertName(t.name.value)}Schema()`).join(', ');
+      const unionElementsCount = node.types?.length ?? 0;
+
+      const union =
+        unionElementsCount > 1 ? indent(`return z.union([${unionElements}])`) : indent(`return ${unionElements}`);
+
+      const result = new DeclarationBlock({})
+        .export()
+        .asKind('function')
+        .withName(`${unionName}Schema()`)
+        .withBlock(union);
+
+      return result.string;
+    },
   };
 };
 
@@ -192,6 +209,11 @@ const generateNameNodeZodSchema = (
     return `${enumName}Schema`;
   }
 
+  if (typ?.astNode?.kind === 'UnionTypeDefinition') {
+    const enumName = tsVisitor.convertName(typ.astNode.name.value);
+    return `${enumName}Schema()`;
+  }
+
   return zod4Scalar(config, tsVisitor, node.value);
 };
 
diff --git a/tests/myzod.spec.ts b/tests/myzod.spec.ts
@@ -538,6 +538,104 @@ describe('myzod', () => {
         expect(result.content).not.toContain(wantNotContain);
       }
     });
+
+    it('generate union types', async () => {
+      const schema = buildSchema(/* GraphQL */ `
+        type Square {
+          size: Int
+        }
+        type Circle {
+          radius: Int
+        }
+        union Shape = Circle | Square
+      `);
+
+      const result = await plugin(
+        schema,
+        [],
+        {
+          schema: 'myzod',
+          withObjectType: true,
+        },
+        {}
+      );
+
+      const wantContains = [
+        // Shape Schema
+        'export function ShapeSchema() {',
+        'myzod.union([CircleSchema(), SquareSchema()])',
+        '}',
+      ];
+      for (const wantContain of wantContains) {
+        expect(result.content).toContain(wantContain);
+      }
+    });
+
+    it('generate union types with single element', async () => {
+      const schema = buildSchema(/* GraphQL */ `
+        type Square {
+          size: Int
+        }
+        type Circle {
+          radius: Int
+        }
+        union Shape = Circle | Square
+
+        type Geometry {
+          shape: Shape
+        }
+      `);
+
+      const result = await plugin(
+        schema,
+        [],
+        {
+          schema: 'myzod',
+          withObjectType: true,
+        },
+        {}
+      );
+
+      const wantContains = [
+        'export function GeometrySchema(): myzod.Type<Geometry> {',
+        'return myzod.object({',
+        "__typename: myzod.literal('Geometry').optional(),",
+        'shape: ShapeSchema().optional().nullable()',
+        '}',
+      ];
+      for (const wantContain of wantContains) {
+        expect(result.content).toContain(wantContain);
+      }
+    });
+
+    it('correctly reference generated union types', async () => {
+      const schema = buildSchema(/* GraphQL */ `
+        type Circle {
+          radius: Int
+        }
+        union Shape = Circle
+      `);
+
+      const result = await plugin(
+        schema,
+        [],
+        {
+          schema: 'myzod',
+          withObjectType: true,
+        },
+        {}
+      );
+
+      const wantContains = [
+        // Shape Schema
+        'export function ShapeSchema() {',
+        'CircleSchema()',
+        '}',
+      ];
+      for (const wantContain of wantContains) {
+        expect(result.content).toContain(wantContain);
+      }
+    });
   });
 
   it('properly generates custom directive values', async () => {
diff --git a/tests/yup.spec.ts b/tests/yup.spec.ts
@@ -451,6 +451,104 @@ describe('yup', () => {
         expect(result.content).not.toContain(wantNotContain);
       }
     });
+
+    it('generate union types', async () => {
+      const schema = buildSchema(/* GraphQL */ `
+        type Square {
+          size: Int
+        }
+        type Circle {
+          radius: Int
+        }
+        union Shape = Circle | Square
+      `);
+
+      const result = await plugin(
+        schema,
+        [],
+        {
+          schema: 'yup',
+          withObjectType: true,
+        },
+        {}
+      );
+
+      const wantContains = [
+        // Shape Schema
+        'export function ShapeSchema(): yup.BaseSchema<Shape> {',
+        'union<Shape>(CircleSchema(), SquareSchema())',
+        '}',
+      ];
+      for (const wantContain of wantContains) {
+        expect(result.content).toContain(wantContain);
+      }
+    });
+
+    it('generate union types with single element', async () => {
+      const schema = buildSchema(/* GraphQL */ `
+        type Square {
+          size: Int
+        }
+        type Circle {
+          radius: Int
+        }
+        union Shape = Circle | Square
+
+        type Geometry {
+          shape: Shape
+        }
+      `);
+
+      const result = await plugin(
+        schema,
+        [],
+        {
+          schema: 'yup',
+          withObjectType: true,
+        },
+        {}
+      );
+
+      const wantContains = [
+        'export function GeometrySchema(): yup.SchemaOf<Geometry> {',
+        'return yup.object({',
+        "__typename: yup.mixed().oneOf(['Geometry', undefined]),",
+        'shape: yup.mixed()',
+        '})',
+      ];
+      for (const wantContain of wantContains) {
+        expect(result.content).toContain(wantContain);
+      }
+    });
+
+    it('correctly reference generated union types', async () => {
+      const schema = buildSchema(/* GraphQL */ `
+        type Circle {
+          radius: Int
+        }
+        union Shape = Circle
+      `);
+
+      const result = await plugin(
+        schema,
+        [],
+        {
+          schema: 'yup',
+          withObjectType: true,
+        },
+        {}
+      );
+
+      const wantContains = [
+        // Shape Schema
+        'export function ShapeSchema(): yup.BaseSchema<Shape> {',
+        'CircleSchema()',
+        '}',
+      ];
+      for (const wantContain of wantContains) {
+        expect(result.content).toContain(wantContain);
+      }
+    });
   });
 
   it('properly generates custom directive values', async () => {
diff --git a/tests/zod.spec.ts b/tests/zod.spec.ts
