fix: adding support for imported QualifiedName (#243)

* test: adding 2 failing cases

* feat: adding type identification in type traversal

* fix: first test now passes (external imports)

* feat: excluding already referenced types

* adding more tests

* fix: typo

src/core/generate.test.ts

@@ -929,6 +929,38 @@ describe("generate", () => {
       `);
       });
     });
+
+    describe("reference to an enum value via external import Qualified Name X.Y", () => {
+      const input = "./person";
+
+      const sourceText = `
+        import { PersonEnum } from "${input}"
+        
+        export interface Hero {
+          id: number
+          hero: PersonEnum.Hero
+        }
+        `;
+
+      const { getZodSchemasFile } = generate({
+        sourceText,
+      });
+
+      it("should generate the zod schemas with right import", () => {
+        expect(getZodSchemasFile(input)).toMatchInlineSnapshot(`
+          "// Generated by ts-to-zod
+          import { z } from "zod";
+          
+          import { PersonEnum } from "${input}";
+          
+          export const heroSchema = z.object({
+              id: z.number(),
+              hero: z.literal(PersonEnum.Hero)
+          });
+          "
+        `);
+      });
+    });
   });
 
   describe("with input/output mappings to manage imports", () => {
@@ -1161,5 +1193,84 @@ describe("generate", () => {
       `);
       });
     });
+
+    describe("reference to an enum value via import Qualified Name X.Y referenced in mapping", () => {
+      const input = "./person";
+      const output = "./person.zod";
+      const inputOutputMappings: InputOutputMapping[] = [{ input, output }];
+
+      const sourceText = `
+        import { PersonEnum } from "${input}"
+        
+        export interface Hero {
+          id: number
+          hero: PersonEnum.Hero
+        }
+        `;
+
+      const { getZodSchemasFile } = generate({
+        sourceText,
+        inputOutputMappings,
+      });
+
+      it("should generate the zod schemas with right import", () => {
+        expect(getZodSchemasFile(input)).toMatchInlineSnapshot(`
+          "// Generated by ts-to-zod
+          import { z } from "zod";
+          
+          import { PersonEnum } from "${input}";
+          
+          export const heroSchema = z.object({
+              id: z.number(),
+              hero: z.literal(PersonEnum.Hero)
+          });
+          "
+        `);
+      });
+    });
+
+    describe("with useless imports", () => {
+      const input = "./hero";
+      const output = "./hero.zod";
+
+      const input2 = "./superhero";
+      const output2 = "./superhero.zod";
+      const inputOutputMappings: InputOutputMapping[] = [
+        { input, output },
+        { input: input2, output: output2 },
+      ];
+
+      const sourceText = `
+      import { Hero, Villain } from "${input}"
+      import { SuperHero } from "${input2}"
+    
+      export interface Person {
+        id: number
+        hero: Hero
+        villain: Villain
+      }
+      `;
+
+      const { getZodSchemasFile } = generate({
+        sourceText,
+        inputOutputMappings,
+      });
+
+      it("should generate the zod schemas with right import", () => {
+        expect(getZodSchemasFile(input)).toMatchInlineSnapshot(`
+        "// Generated by ts-to-zod
+        import { z } from "zod";
+        
+        import { heroSchema, villainSchema } from "${output}";
+        
+        export const personSchema = z.object({
+            id: z.number(),
+            hero: heroSchema,
+            villain: villainSchema
+        });
+        "
+      `);
+      });
+    });
   });
 });

src/core/generate.ts

@@ -10,8 +10,9 @@ import {
 import { getSimplifiedJsDocTags } from "../utils/getSimplifiedJsDocTags";
 import { resolveModules } from "../utils/resolveModules";
 import {
-  getExtractedTypeNames,
+  getReferencedTypeNames,
   isTypeNode,
+  TypeNameReference,
   TypeNode,
 } from "../utils/traverseTypes";
 
@@ -103,21 +104,33 @@ export function generate({
   const nodes: Array<TypeNode> = [];
 
   // declare a map to store the interface name and its corresponding zod schema
-  const typeNameMapping = new Map<string, TypeNode>();
+  const typeNameMapping = new Map<string, TypeNode | ts.ImportDeclaration>();
 
-  // handling existing import statements
-  const importNamesAvailable = new Set<string>();
-  const importNodes: ts.ImportDeclaration[] = [];
-  const importNamesUsed: string[] = [];
+  /**
+   * Following const are keeping track of all the things import-related
+   */
+  // All import nodes in the source file
+  const zodImportNodes: ts.ImportDeclaration[] = [];
+
+  // Keep track of all the external import names available in the source file
+  const externalImportNamesAvailable = new Set<string>();
+
+  // Keep track of all the imports that have an entry in the config file
   const importedZodNamesAvailable = new Map<string, string>();
 
-  const typesNeedToBeExtracted = new Set<string>();
+  // Keep track of all referenced types in the source file
+  const candidateTypesToBeExtracted = new Set<TypeNameReference>();
 
   const typeNameMapBuilder = (node: ts.Node) => {
     if (isTypeNode(node)) {
       typeNameMapping.set(node.name.text, node);
+      return;
     }
+
     if (ts.isImportDeclaration(node) && node.importClause) {
+      const identifiers = getImportIdentifiers(node);
+      identifiers.forEach((i) => typeNameMapping.set(i, node));
+
       // Check if we're importing from a mapped file
       const eligibleMapping = inputOutputMappings.find(
         (io: InputOutputMapping) =>
@@ -136,16 +149,15 @@ export function generate({
           importedZodNamesAvailable.set(i, schemaMethod(i))
         );
 
-        const importNode = createImportNode(
+        const zodImportNode = createImportNode(
           identifiers.map(schemaMethod),
           eligibleMapping.output
         );
-        importNodes.push(importNode); // We assume all identifiers will be used so pushing the whole node
+        zodImportNodes.push(zodImportNode);
       }
       // Not a Zod import, handling it as 3rd party import later on
       else {
-        const identifiers = getImportIdentifiers(node);
-        identifiers.forEach((i) => importNamesAvailable.add(i));
+        identifiers.forEach((i) => externalImportNamesAvailable.add(i));
       }
     }
   };
@@ -162,34 +174,66 @@ export function generate({
       if (!jsDocTagFilter(tags)) return;
       if (!nameFilter(node.name.text)) return;
 
-      const typeNames = getExtractedTypeNames(node, sourceFile);
-      typeNames.forEach((typeName) => {
-        typesNeedToBeExtracted.add(typeName);
+      const typeNames = getReferencedTypeNames(node, sourceFile);
+      typeNames.forEach((typeRef) => {
+        candidateTypesToBeExtracted.add(typeRef);
       });
     }
   };
   ts.forEachChild(sourceFile, visitor);
 
+  // All external import names actually used in the source file
+  const importNamesUsed: string[] = [];
+
+  // All zod imports actually used in the source file
   const importedZodSchemas = new Set<string>();
 
-  typesNeedToBeExtracted.forEach((typeName) => {
-    const node = typeNameMapping.get(typeName);
+  // All original import to keep in the target
+  const importsToKeep = new Map<ts.ImportDeclaration, string[]>();
+
+  /**
+   * We browse all the extracted type references from the source file
+   * To check if they reference a node from the file or if they are imported
+   */
+  candidateTypesToBeExtracted.forEach((typeRef) => {
+    const node = typeNameMapping.get(typeRef.name);
+
     if (node) {
-      nodes.push(node);
-      return;
+      // If we have a reference in the file, we add it to the nodes, no import needed
+      if (isTypeNode(node)) {
+        nodes.push(node);
+        return;
+      }
+
+      // If the reference is part of a qualified name, we need to import it from the same file
+      if (typeRef.partOfQualifiedName) {
+        const identifiers = importsToKeep.get(node);
+        if (identifiers) {
+          identifiers.push(typeRef.name);
+        } else {
+          importsToKeep.set(node, [typeRef.name]);
+        }
+        return;
+      }
     }
-    if (importNamesAvailable.has(typeName)) {
-      importNamesUsed.push(typeName);
+
+    // If the reference is coming from an external import, we'll need to generate a specific statement
+    // and keep the external import
+    if (externalImportNamesAvailable.has(typeRef.name)) {
+      importNamesUsed.push(typeRef.name);
       return;
     }
-    if (importedZodNamesAvailable.has(typeName)) {
-      importedZodSchemas.add(importedZodNamesAvailable.get(typeName) as string);
+
+    // If the reference is coming from a mapped import, we'll import the corresponding zod schema
+    if (importedZodNamesAvailable.has(typeRef.name)) {
+      importedZodSchemas.add(
+        importedZodNamesAvailable.get(typeRef.name) as string
+      );
       return;
     }
   });
 
   // Generate zod schemas for type nodes
-
   const getDependencyName = (identifierName: string) => {
     if (importedZodNamesAvailable.has(identifierName)) {
       return importedZodNamesAvailable.get(identifierName) as string;
@@ -237,7 +281,9 @@ export function generate({
     string,
     { typeName: string; value: ts.VariableStatement }
   >();
-  const typeImports: Set<string> = new Set();
+
+  // Keep track of types which need to be imported from the source file
+  const sourceTypeImports: 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>();
@@ -266,14 +312,14 @@ export function generate({
           if (notGeneratedDependencies.length === 0) {
             done = false;
             if (isCircular) {
-              typeImports.add(typeName);
+              sourceTypeImports.add(typeName);
               statements.set(varName, {
                 value: transformRecursiveSchema("z", statement, typeName),
                 typeName,
               });
             } else {
               if (requiresImport) {
-                typeImports.add(typeName);
+                sourceTypeImports.add(typeName);
               }
               statements.set(varName, { value: statement, typeName });
             }
@@ -300,7 +346,7 @@ export function generate({
         !zodSchemasWithMissingDependencies.has(varName)
     )
     .forEach(({ varName, statement, typeName }) => {
-      typeImports.add(typeName);
+      sourceTypeImports.add(typeName);
       statements.set(varName, {
         value: transformRecursiveSchema("z", statement, typeName),
         typeName,
@@ -332,21 +378,39 @@ ${Array.from(zodSchemasWithMissingDependencies).join("\n")}`
 
   const transformedSourceText = printerWithComments.printFile(sourceFile);
 
-  const typeImportsValues = Array.from(typeImports.values());
+  const zodImportToOutput = zodImportNodes.filter((node) => {
+    const nodeIdentifiers = getImportIdentifiers(node);
+    return nodeIdentifiers.some((i) => importedZodSchemas.has(i));
+  });
+
+  const originalImportsToOutput = Array.from(importsToKeep.keys()).map((node) =>
+    createImportNode(
+      importsToKeep.get(node)!,
+      (node.moduleSpecifier as ts.StringLiteral).text
+    )
+  );
+
+  const sourceTypeImportsValues = Array.from(sourceTypeImports.values());
   const getZodSchemasFile = (
     typesImportPath: string
   ) => `// Generated by ts-to-zod
 import { z } from "zod";
 ${
-  typeImportsValues.length
-    ? `import { ${typeImportsValues.join(", ")} } from "${typesImportPath}";\n`
+  sourceTypeImportsValues.length
+    ? `import { ${sourceTypeImportsValues.join(
+        ", "
+      )} } from "${typesImportPath}";\n`
     : ""
 }
 ${
-  importNodes.length
-    ? importNodes.map((node) => print(node)).join("\n") + "\n\n"
+  zodImportToOutput.length
+    ? zodImportToOutput.map((node) => print(node)).join("\n") + "\n\n"
     : ""
-}${Array.from(statements.values())
+}${
+    originalImportsToOutput.length
+      ? originalImportsToOutput.map((node) => print(node)).join("\n") + "\n\n"
+      : ""
+  }${Array.from(statements.values())
     .map((statement) => print(statement.value))
     .join("\n\n")}
 `;
@@ -448,7 +512,7 @@ ${Array.from(statements.values())
     /**
      * `true` if zodSchemaFile have some resolvable circular dependencies
      */
-    hasCircularDependencies: typeImportsValues.length > 0,
+    hasCircularDependencies: sourceTypeImportsValues.length > 0,
   };
 }