Adds support for use of enum types as literals and nativeEnums (#40)

* Checks in test documenting behavior * Implements enum import & declartion * Updates example to include enum * Updates tests to document expected case * Add handler for enum literals * Updates example * Addresses #38 (comment) * Addresses #38 (comment) * Addresses #40 (comment)
fabien0102 · Jul 29, 2021 · 8bf26ed · 8bf26ed
1 parent 92f1d00
commit 8bf26ed
Show file tree

Hide file tree

Showing 6 changed files with 168 additions and 25 deletions.
diff --git a/example/heros.ts b/example/heros.ts
@@ -1,6 +1,16 @@
+export enum EnemyPower {
+  Flight = "flight",
+  Strength = "strength",
+  Speed = "speed",
+}
+
+export type SpeedEnemy = {
+  power: EnemyPower.Speed;
+};
+
 export interface Enemy {
   name: string;
-  powers: string[];
+  powers: EnemyPower[];
   inPrison: boolean;
 }
 
@@ -13,7 +23,7 @@ export interface Superman {
 
 export interface Villain {
   name: string;
-  powers: string[];
+  powers: EnemyPower[];
   friends: Villain[];
   canBeTrusted: never;
 }

diff --git a/example/heros.zod.ts b/example/heros.zod.ts
@@ -1,10 +1,16 @@
 // Generated by ts-to-zod
 import { z } from "zod";
-import { Villain } from "./heros";
+import { EnemyPower, Villain } from "./heros";
+
+export const enemyPowerSchema = z.nativeEnum(EnemyPower);
+
+export const speedEnemySchema = z.object({
+  power: z.literal(EnemyPower.Speed),
+});
 
 export const enemySchema = z.object({
   name: z.string(),
-  powers: z.array(z.string()),
+  powers: z.array(enemyPowerSchema),
   inPrison: z.boolean(),
 });
 
@@ -22,7 +28,7 @@ export const supermanSchema = z.object({
 export const villainSchema: z.ZodSchema<Villain> = z.lazy(() =>
   z.object({
     name: z.string(),
-    powers: z.array(z.string()),
+    powers: z.array(enemyPowerSchema),
     friends: z.array(villainSchema),
     canBeTrusted: z.never(),
   })

diff --git a/src/core/generate.test.ts b/src/core/generate.test.ts
@@ -72,6 +72,66 @@ describe("generate", () => {
         "
       `);
     });
+    it("should not have any errors", () => {
+      expect(errors.length).toBe(0);
+    });
+  });
+
+  describe("with enums", () => {
+    const sourceText = `
+      export enum Superhero { 
+        Superman = "superman"
+        ClarkKent = "clark-kent"
+      };
+
+      export type FavoriteSuperhero = {
+        superhero: Superhero.Superman
+      };
+      `;
+
+    const { getZodSchemasFile, getIntegrationTestFile, errors } = generate({
+      sourceText,
+    });
+
+    it("should generate the zod schemas", () => {
+      expect(getZodSchemasFile("./superhero")).toMatchInlineSnapshot(`
+        "// Generated by ts-to-zod
+        import { z } from \\"zod\\";
+        import { Superhero } from \\"./superhero\\";
+
+        export const superheroSchema = z.nativeEnum(Superhero);
+
+        export const favoriteSuperheroSchema = z.object({
+            superhero: z.literal(Superhero.Superman)
+        });
+        "
+      `);
+    });
+
+    it("should generate the integration tests", () => {
+      expect(getIntegrationTestFile("./superhero", "superhero.zod"))
+        .toMatchInlineSnapshot(`
+        "// Generated by ts-to-zod
+        import { z } from \\"zod\\";
+
+        import * as spec from \\"./superhero\\";
+        import * as generated from \\"superhero.zod\\";
+
+        // eslint-disable-next-line @typescript-eslint/no-unused-vars
+        function expectType<T>(_: T) {
+          /* noop */
+        }
+
+        export type superheroSchemaInferredType = z.infer<typeof generated.superheroSchema>;
+
+        export type favoriteSuperheroSchemaInferredType = z.infer<typeof generated.favoriteSuperheroSchema>;
+        expectType<spec.Superhero>({} as superheroSchemaInferredType)
+        expectType<superheroSchemaInferredType>({} as spec.Superhero)
+        expectType<spec.FavoriteSuperhero>({} as favoriteSuperheroSchemaInferredType)
+        expectType<favoriteSuperheroSchemaInferredType>({} as spec.FavoriteSuperhero)
+        "
+      `);
+    });
 
     it("should not have any errors", () => {
       expect(errors.length).toBe(0);

diff --git a/src/core/generate.ts b/src/core/generate.ts
@@ -53,10 +53,16 @@ export function generate({
   );
 
   // Extract the nodes (interface declarations & type aliases)
-  const nodes: Array<ts.InterfaceDeclaration | ts.TypeAliasDeclaration> = [];
+  const nodes: Array<
+    ts.InterfaceDeclaration | ts.TypeAliasDeclaration | ts.EnumDeclaration
+  > = [];
 
   const visitor = (node: ts.Node) => {
-    if (ts.isInterfaceDeclaration(node) || ts.isTypeAliasDeclaration(node)) {
+    if (
+      ts.isInterfaceDeclaration(node) ||
+      ts.isTypeAliasDeclaration(node) ||
+      ts.isEnumDeclaration(node)
+    ) {
       if (nameFilter(node.name.text)) {
         nodes.push(node);
       }
@@ -91,23 +97,28 @@ export function generate({
   while (statements.size !== zodSchemas.length && n < maxRun) {
     zodSchemas
       .filter(({ varName }) => !statements.has(varName))
-      .forEach(({ varName, dependencies, statement, typeName }) => {
-        const isCircular = dependencies.includes(varName);
-        const missingDependencies = dependencies
-          .filter((dep) => dep !== varName)
-          .filter((dep) => !statements.has(dep));
-        if (missingDependencies.length === 0) {
-          if (isCircular) {
-            typeImports.add(typeName);
-            statements.set(varName, {
-              value: transformRecursiveSchema("z", statement, typeName),
-              typeName,
-            });
-          } else {
-            statements.set(varName, { value: statement, typeName });
+      .forEach(
+        ({ varName, dependencies, statement, typeName, requiresImport }) => {
+          const isCircular = dependencies.includes(varName);
+          const missingDependencies = dependencies
+            .filter((dep) => dep !== varName)
+            .filter((dep) => !statements.has(dep));
+          if (missingDependencies.length === 0) {
+            if (isCircular) {
+              typeImports.add(typeName);
+              statements.set(varName, {
+                value: transformRecursiveSchema("z", statement, typeName),
+                typeName,
+              });
+            } else {
+              if (requiresImport) {
+                typeImports.add(typeName);
+              }
+              statements.set(varName, { value: statement, typeName });
+            }
           }
         }
-      });
+      );
 
     n++; // Just a safety net to avoid infinity loops
   }

diff --git a/src/core/generateZodSchema.test.ts b/src/core/generateZodSchema.test.ts
@@ -95,6 +95,29 @@ describe("generateZodSchema", () => {
     );
   });
 
+  it("should generate a literal schema (enum)", () => {
+    const source = `
+    export type BestSuperhero = {
+      superhero: Superhero.Superman
+    };
+    `;
+    expect(generate(source)).toMatchInlineSnapshot(`
+      "export const bestSuperheroSchema = z.object({
+          superhero: z.literal(Superhero.Superman)
+      });"
+    `);
+  });
+
+  it("should generate a nativeEnum schema", () => {
+    const source = `export enum Superhero = { 
+      Superman = "superman",
+      ClarkKent = "clark_kent",
+    };`;
+    expect(generate(source)).toMatchInlineSnapshot(
+      `"export const superheroSchema = z.nativeEnum(Superhero);"`
+    );
+  });
+
   it("should generate a never", () => {
     const source = `export type CanBeatZod = never;`;
     expect(generate(source)).toMatchInlineSnapshot(
@@ -627,8 +650,15 @@ function generate(sourceText: string, z?: string) {
   );
   const declaration = findNode(
     sourceFile,
-    (node): node is ts.InterfaceDeclaration | ts.TypeAliasDeclaration =>
-      ts.isInterfaceDeclaration(node) || ts.isTypeAliasDeclaration(node)
+    (
+      node
+    ): node is
+      | ts.InterfaceDeclaration
+      | ts.TypeAliasDeclaration
+      | ts.EnumDeclaration =>
+      ts.isInterfaceDeclaration(node) ||
+      ts.isTypeAliasDeclaration(node) ||
+      ts.isEnumDeclaration(node)
   );
   if (!declaration) {
     throw new Error("No `type` or `interface` found!");

diff --git a/src/core/generateZodSchema.ts b/src/core/generateZodSchema.ts
@@ -19,7 +19,7 @@ export interface GenerateZodSchemaProps {
   /**
    * Interface or type node
    */
-  node: ts.InterfaceDeclaration | ts.TypeAliasDeclaration;
+  node: ts.InterfaceDeclaration | ts.TypeAliasDeclaration | ts.EnumDeclaration;
 
   /**
    * Zod import value.
@@ -57,6 +57,7 @@ export function generateZodSchemaVariableStatement({
 }: GenerateZodSchemaProps) {
   let schema: ts.CallExpression | ts.Identifier | undefined;
   const dependencies: string[] = [];
+  let requiresImport = false;
 
   if (ts.isInterfaceDeclaration(node)) {
     let baseSchema: string | undefined;
@@ -100,6 +101,11 @@ export function generateZodSchemaVariableStatement({
     });
   }
 
+  if (ts.isEnumDeclaration(node)) {
+    schema = buildZodSchema(zodImportValue, "nativeEnum", [node.name]);
+    requiresImport = true;
+  }
+
   return {
     dependencies: uniq(dependencies),
     statement: f.createVariableStatement(
@@ -116,6 +122,7 @@ export function generateZodSchemaVariableStatement({
         ts.NodeFlags.Const
       )
     ),
+    requiresImport,
   };
 }
 
@@ -452,6 +459,25 @@ function buildZodPrimitive({
     }
   }
 
+  // Deal with enums used as literals
+  if (
+    ts.isTypeReferenceNode(typeNode) &&
+    ts.isQualifiedName(typeNode.typeName) &&
+    ts.isIdentifier(typeNode.typeName.left)
+  ) {
+    return buildZodSchema(
+      z,
+      "literal",
+      [
+        f.createPropertyAccessExpression(
+          typeNode.typeName.left,
+          typeNode.typeName.right
+        ),
+      ],
+      zodProperties
+    );
+  }
+
   if (ts.isArrayTypeNode(typeNode)) {
     return buildZodSchema(
       z,