feat!: Inheritance and reference type search for name filtering (#104)

* feat: name filtering include heritage & properties reference type

* test: add test case for type search

* feat: update traverse type defined

Co-authored-by: Fabien BERNARD <fabien0102@hotmail.com>

* Extract isTypeNode utils

Co-authored-by: Fabien BERNARD <fabien0102@hotmail.com>
Co-authored-by: Fabien BERNARD <fabien0102@gmail.com>

src/core/generate.test.ts

@@ -247,6 +247,35 @@ describe("generate", () => {
     });
   });
 
+  describe("inheritance and reference type search", () => {
+    const sourceText = `
+    export type Name = "superman" | "clark kent" | "kal-l";
+    export interface Superman {
+      name: Name;
+    }`;
+
+    const { getZodSchemasFile } = generate({
+      sourceText,
+      nameFilter: (id) => id === "Superman",
+      getSchemaName: (id) => id.toLowerCase(),
+      keepComments: true,
+    });
+
+    it("should generate superman schema", () => {
+      expect(getZodSchemasFile("./hero")).toMatchInlineSnapshot(`
+        "// Generated by ts-to-zod
+        import { z } from \\"zod\\";
+
+        export const name = z.union([z.literal(\\"superman\\"), z.literal(\\"clark kent\\"), z.literal(\\"kal-l\\")]);
+
+        export const superman = z.object({
+            name: name
+        });
+        "
+      `);
+    });
+  });
+
   describe("with jsdocTags filter", () => {
     it("should generate only types with @zod", () => {
       const sourceText = `

src/core/generate.ts

@@ -4,6 +4,11 @@ import ts from "typescript";
 import { JSDocTagFilter, NameFilter } from "../config";
 import { getSimplifiedJsDocTags } from "../utils/getSimplifiedJsDocTags";
 import { resolveModules } from "../utils/resolveModules";
+import {
+  getExtractedTypeNames,
+  isTypeNode,
+  TypeNode,
+} from "../utils/traverseTypes";
 import { generateIntegrationTests } from "./generateIntegrationTests";
 import { generateZodInferredType } from "./generateZodInferredType";
 import { generateZodSchemaVariableStatement } from "./generateZodSchema";
@@ -67,10 +72,19 @@ export function generate({
   const sourceFile = resolveModules(sourceText);
 
   // Extract the nodes (interface declarations & type aliases)
-  const nodes: Array<
-    ts.InterfaceDeclaration | ts.TypeAliasDeclaration | ts.EnumDeclaration
-  > = [];
+  const nodes: Array<TypeNode> = [];
 
+  // declare a map to store the interface name and its corresponding zod schema
+  const typeNameMapping = new Map<string, TypeNode>();
+
+  const typesNeedToBeExtracted = new Set<string>();
+
+  const typeNameMapBuilder = (node: ts.Node) => {
+    if (isTypeNode(node)) {
+      typeNameMapping.set(node.name.text, node);
+    }
+  };
+  ts.forEachChild(sourceFile, typeNameMapBuilder);
   const visitor = (node: ts.Node) => {
     if (
       ts.isInterfaceDeclaration(node) ||
@@ -81,11 +95,26 @@ export function generate({
       const tags = getSimplifiedJsDocTags(jsDoc);
       if (!jsDocTagFilter(tags)) return;
       if (!nameFilter(node.name.text)) return;
-      nodes.push(node);
+
+      const typeNames = getExtractedTypeNames(
+        node,
+        sourceFile,
+        typeNameMapping
+      );
+      typeNames.forEach((typeName) => {
+        typesNeedToBeExtracted.add(typeName);
+      });
     }
   };
   ts.forEachChild(sourceFile, visitor);
 
+  typesNeedToBeExtracted.forEach((typeName) => {
+    const node = typeNameMapping.get(typeName);
+    if (node) {
+      nodes.push(node);
+    }
+  });
+
   // Generate zod schemas
   const zodSchemas = nodes.map((node) => {
     const typeName = node.name.text;

src/utils/traverseTypes.ts

@@ -0,0 +1,71 @@
+import ts from "typescript";
+
+export type TypeNode = (
+  | ts.InterfaceDeclaration
+  | ts.TypeAliasDeclaration
+  | ts.EnumDeclaration
+) & { visited?: boolean };
+
+export function isTypeNode(node: ts.Node): node is TypeNode {
+  return (
+    ts.isInterfaceDeclaration(node) ||
+    ts.isTypeAliasDeclaration(node) ||
+    ts.isEnumDeclaration(node)
+  );
+}
+
+export function getExtractedTypeNames(
+  node: TypeNode,
+  sourceFile: ts.SourceFile,
+  typeNameMapping: Map<string, TypeNode>
+) {
+  const referenceTypeNames: string[] = [];
+
+  const recursiveExtract = (node: TypeNode) => {
+    if (node.visited) {
+      return;
+    }
+
+    const heritageClauses = (node as ts.InterfaceDeclaration).heritageClauses;
+
+    if (heritageClauses) {
+      heritageClauses.forEach((clause) => {
+        const extensionTypes = clause.types;
+        extensionTypes.forEach((extensionTypeNode) => {
+          const typeName = extensionTypeNode.expression.getText(sourceFile);
+          const typeNode = typeNameMapping.get(typeName);
+
+          referenceTypeNames.push(typeName);
+
+          if (typeNode) {
+            typeNode.visited = true;
+            recursiveExtract(typeNode);
+          }
+        });
+      });
+    }
+
+    node.forEachChild((child) => {
+      const childNode = child as ts.PropertySignature;
+      if (childNode.kind !== ts.SyntaxKind.PropertySignature) {
+        return;
+      }
+
+      if (childNode.type?.kind === ts.SyntaxKind.TypeReference) {
+        const typeNode = typeNameMapping.get(
+          childNode.type.getText(sourceFile)
+        );
+
+        referenceTypeNames.push(childNode.type.getText(sourceFile));
+
+        if (typeNode) {
+          typeNode.visited = true;
+          recursiveExtract(typeNode);
+        }
+      }
+    });
+  };
+
+  recursiveExtract(node);
+  return [node.name.text, ...referenceTypeNames];
+}