feat: add type notation to type imports (#264)

fabien0102 · Aug 29, 2024 · 941169d · 941169d
1 parent 38c795f
commit 941169d
Show file tree

Hide file tree

Showing 3 changed files with 80 additions and 38 deletions.
diff --git a/src/core/generate.test.ts b/src/core/generate.test.ts
@@ -326,7 +326,7 @@ describe("generate", () => {
       expect(getZodSchemasFile("./villain")).toMatchInlineSnapshot(`
         "// Generated by ts-to-zod
         import { z } from "zod";
-        import { Villain, EvilPlan, EvilPlanDetails } from "./villain";
+        import { type Villain, type EvilPlan, type EvilPlanDetails } from "./villain";
 
         export const villainSchema: z.ZodSchema<Villain> = z.lazy(() => z.object({
             name: z.string(),
@@ -961,6 +961,41 @@ describe("generate", () => {
         `);
       });
     });
+
+    describe("with mixed imports", () => {
+      const input = "./person";
+
+      const sourceText = `
+        import { PersonEnum } from "${input}"
+
+        export interface Hero {
+          id: number
+          hero: PersonEnum.Hero
+          parent: Hero
+        }
+        `;
+
+      const { getZodSchemasFile } = generate({
+        sourceText,
+      });
+
+      it("should add type notation to non-enum imports", () => {
+        expect(getZodSchemasFile(input)).toMatchInlineSnapshot(`
+          "// Generated by ts-to-zod
+          import { z } from "zod";
+          import { type Hero } from "${input}";
+
+          import { PersonEnum } from "${input}";
+
+          export const heroSchema: z.ZodSchema<Hero> = z.lazy(() => z.object({
+              id: z.number(),
+              hero: z.literal(PersonEnum.Hero),
+              parent: heroSchema
+          }));
+          "
+        `);
+      });
+    });
   });
 
   describe("with input/output mappings to manage imports", () => {

diff --git a/src/core/generate.ts b/src/core/generate.ts
@@ -266,7 +266,7 @@ export function generate({
         varName,
         zodImportValue: "z",
       }),
-      requiresImport: false,
+      enumImport: false,
       typeName: importName,
       varName,
     };
@@ -282,8 +282,9 @@ export function generate({
     { typeName: string; value: ts.VariableStatement }
   >();
 
-  // Keep track of types which need to be imported from the source file
+  // Keep track of types/enums which need to be imported from the source file
   const sourceTypeImports: Set<string> = new Set();
+  const sourceEnumImports: Set<string> = new Set();
 
   // Zod schemas with direct or indirect dependencies that are not in `zodSchemas`, won't be generated
   const zodSchemasWithMissingDependencies = new Set<string>();
@@ -302,40 +303,38 @@ export function generate({
           !statements.has(varName) &&
           !zodSchemasWithMissingDependencies.has(varName)
       )
-      .forEach(
-        ({ varName, dependencies, statement, typeName, requiresImport }) => {
-          const isCircular = dependencies.includes(varName);
-          const notGeneratedDependencies = dependencies
-            .filter((dep) => dep !== varName)
-            .filter((dep) => !statements.has(dep))
-            .filter((dep) => !importedZodSchemas.has(dep));
-          if (notGeneratedDependencies.length === 0) {
-            done = false;
-            if (isCircular) {
-              sourceTypeImports.add(typeName);
-              statements.set(varName, {
-                value: transformRecursiveSchema("z", statement, typeName),
-                typeName,
-              });
-            } else {
-              if (requiresImport) {
-                sourceTypeImports.add(typeName);
-              }
-              statements.set(varName, { value: statement, typeName });
+      .forEach(({ varName, dependencies, statement, typeName, enumImport }) => {
+        const isCircular = dependencies.includes(varName);
+        const notGeneratedDependencies = dependencies
+          .filter((dep) => dep !== varName)
+          .filter((dep) => !statements.has(dep))
+          .filter((dep) => !importedZodSchemas.has(dep));
+        if (notGeneratedDependencies.length === 0) {
+          done = false;
+          if (isCircular) {
+            sourceTypeImports.add(typeName);
+            statements.set(varName, {
+              value: transformRecursiveSchema("z", statement, typeName),
+              typeName,
+            });
+          } else {
+            if (enumImport) {
+              sourceEnumImports.add(typeName);
             }
-          } else if (
-            // Check if every dependency is (in `zodSchemas` and not in `zodSchemasWithMissingDependencies`)
-            !notGeneratedDependencies.every(
-              (dep) =>
-                zodSchemaNames.includes(dep) &&
-                !zodSchemasWithMissingDependencies.has(dep)
-            )
-          ) {
-            done = false;
-            zodSchemasWithMissingDependencies.add(varName);
+            statements.set(varName, { value: statement, typeName });
           }
+        } else if (
+          // Check if every dependency is (in `zodSchemas` and not in `zodSchemasWithMissingDependencies`)
+          !notGeneratedDependencies.every(
+            (dep) =>
+              zodSchemaNames.includes(dep) &&
+              !zodSchemasWithMissingDependencies.has(dep)
+          )
+        ) {
+          done = false;
+          zodSchemasWithMissingDependencies.add(varName);
         }
-      );
+      });
   }
 
   // Generate remaining schemas, which have circular dependencies with loop of length > 1 like: A->B—>C->A
@@ -390,7 +389,15 @@ ${Array.from(zodSchemasWithMissingDependencies).join("\n")}`
     )
   );
 
-  const sourceTypeImportsValues = Array.from(sourceTypeImports.values());
+  const sourceTypeImportsValues = [
+    ...sourceTypeImports.values(),
+    ...sourceEnumImports.values(),
+  ].map((name) => {
+    return sourceEnumImports.has(name)
+      ? name // enum import, no type notation added
+      : `type ${name}`;
+  });
+
   const getZodSchemasFile = (
     typesImportPath: string
   ) => `// Generated by ts-to-zod

diff --git a/src/core/generateZodSchema.ts b/src/core/generateZodSchema.ts
@@ -84,7 +84,7 @@ export function generateZodSchemaVariableStatement({
     | ts.PropertyAccessExpression
     | undefined;
   let dependencies: string[] = [];
-  let requiresImport = false;
+  let enumImport = false;
 
   if (ts.isInterfaceDeclaration(node)) {
     let schemaExtensionClauses: SchemaExtensionClause[] | undefined;
@@ -177,7 +177,7 @@ export function generateZodSchemaVariableStatement({
 
   if (ts.isEnumDeclaration(node)) {
     schema = buildZodSchema(zodImportValue, "nativeEnum", [node.name]);
-    requiresImport = true;
+    enumImport = true;
   }
 
   return {
@@ -196,7 +196,7 @@ export function generateZodSchemaVariableStatement({
         ts.NodeFlags.Const
       )
     ),
-    requiresImport,
+    enumImport,
   };
 }