fix: Various type extractions (#159)

* fix: add support for ParenthesisTypeNodes

* test: adding failing tests

* fix: handling type alias

* fix: typo

* fix: complex array were not traversed

* fix: missing TS helpers in type extractor

src/core/generateZodSchema.test.ts

@@ -1162,6 +1162,23 @@ describe("generateZodSchema", () => {
     `);
   });
 
+  it("should handle parenthesis type nodes", () => {
+    const source = `export interface A {
+      a: (number) | string | null;
+      b: (string)
+      c: (number | string)
+    }
+    `;
+
+    expect(generate(source)).toMatchInlineSnapshot(`
+      "export const aSchema = z.object({
+          a: z.union([z.number(), z.string()]).nullable(),
+          b: z.string(),
+          c: z.union([z.number(), z.string()])
+      });"
+    `);
+  });
+
   it("should allow nullable on optional union properties", () => {
     const source = `export interface A {
       a?: number | string | null;

src/utils/traverseTypes.test.ts

@@ -119,6 +119,140 @@ describe("traverseTypes", () => {
       const result = extractNames(source);
       expect(result).toEqual(["Person", "SuperHero", "Villain"]);
     });
+
+    it("should extract types between parenthesis", () => {
+      const source = `
+        export interface Person {
+            id: number,
+            t: (SuperHero)
+        }`;
+
+      const result = extractNames(source);
+      expect(result).toEqual(["Person", "SuperHero"]);
+    });
+
+    it("should extract union types between parenthesis", () => {
+      const source = `
+        export interface Person {
+            id: number,
+            t: (SuperHero | Villain)
+        }`;
+
+      const result = extractNames(source);
+      expect(result).toEqual(["Person", "SuperHero", "Villain"]);
+    });
+
+    it("should extract type from type alias", () => {
+      const source = `
+        export type Person = SuperHero `;
+
+      const result = extractNames(source);
+      expect(result).toEqual(["Person", "SuperHero"]);
+    });
+
+    it("should extract type from type alias with union", () => {
+      const source = `
+        export type Person = Villain | SuperHero `;
+
+      const result = extractNames(source);
+      expect(result).toEqual(["Person", "Villain", "SuperHero"]);
+    });
+
+    it("should extract type from type alias with array", () => {
+      const source = `
+        export type Person = Villain[] `;
+
+      const result = extractNames(source);
+      expect(result).toEqual(["Person", "Villain"]);
+    });
+
+    it("should extract type from type alias with Array helper", () => {
+      const source = `
+        export type Person = Array<Villain> `;
+
+      const result = extractNames(source);
+      expect(result).toEqual(["Person", "Villain"]);
+    });
+
+    it("should extract type from type alias with Promise helper", () => {
+      const source = `
+        export type Person = Promise<Villain> `;
+
+      const result = extractNames(source);
+      expect(result).toEqual(["Person", "Villain"]);
+    });
+
+    it("should extract type from type alias with Required helper", () => {
+      const source = `
+        export type Person = Required<Villain> `;
+
+      const result = extractNames(source);
+      expect(result).toEqual(["Person", "Villain"]);
+    });
+
+    it("should extract type from type alias with Partial helper", () => {
+      const source = `
+        export type Person = Partial<Villain> `;
+
+      const result = extractNames(source);
+      expect(result).toEqual(["Person", "Villain"]);
+    });
+
+    it("should extract type from type alias with Omit helper", () => {
+      const source = `
+        export type Person = Omit<Villain> `;
+
+      const result = extractNames(source);
+      expect(result).toEqual(["Person", "Villain"]);
+    });
+
+    it("should extract type from type alias with Pick helper", () => {
+      const source = `
+        export type Person = Pick<Villain> `;
+
+      const result = extractNames(source);
+      expect(result).toEqual(["Person", "Villain"]);
+    });
+
+    it("should extract type from type alias with Record helper", () => {
+      const source = `
+        export type Person = Record<string, Villain> `;
+
+      const result = extractNames(source);
+      expect(result).toEqual(["Person", "Villain"]);
+    });
+
+    it("should extract type from type alias with parenthesis", () => {
+      const source = `
+        export type Person = (Villain) `;
+
+      const result = extractNames(source);
+      expect(result).toEqual(["Person", "Villain"]);
+    });
+
+    it("should extract type from type alias with parenthesis", () => {
+      const source = `
+        export type Person = (Villain | Hero)[]`;
+
+      const result = extractNames(source);
+      expect(result).toEqual(["Person", "Villain", "Hero"]);
+    });
+
+    it("should extract type from type alias with object literal", () => {
+      const source = `
+        export type Person = { hero: SuperHero } `;
+
+      const result = extractNames(source);
+      expect(result).toEqual(["Person", "SuperHero"]);
+    });
+
+    it("should extract type from type alias with union & object literal", () => {
+      const source = `
+        export type Person = Villain | { hero: SuperHero } `;
+
+      const result = extractNames(source);
+      expect(result).toEqual(["Person", "Villain", "SuperHero"]);
+    });
   });
 });
 

src/utils/traverseTypes.ts

@@ -1,5 +1,15 @@
 import ts from "typescript";
 
+const typeScriptHelper = [
+  "Array",
+  "Promise",
+  "Omit",
+  "Pick",
+  "Record",
+  "Partial",
+  "Required",
+];
+
 export type TypeNode =
   | ts.InterfaceDeclaration
   | ts.TypeAliasDeclaration
@@ -14,55 +24,75 @@ export function isTypeNode(node: ts.Node): node is TypeNode {
 }
 
 export function getExtractedTypeNames(
-  node: TypeNode,
+  node: ts.InterfaceDeclaration | ts.TypeAliasDeclaration | ts.EnumDeclaration,
   sourceFile: ts.SourceFile
 ): string[] {
   const referenceTypeNames = new Set<string>();
+
+  // Adding the node name
   referenceTypeNames.add(node.name.text);
 
-  const heritageClauses = (node as ts.InterfaceDeclaration).heritageClauses;
+  const visitorExtract = (child: ts.Node) => {
+    if (!ts.isPropertySignature(child)) {
+      return;
+    }
+
+    const childNode = child as ts.PropertySignature;
+    if (childNode.type) {
+      handleTypeNode(childNode.type);
+    }
+  };
 
-  if (heritageClauses) {
-    heritageClauses.forEach((clause) => {
-      const extensionTypes = clause.types;
-      extensionTypes.forEach((extensionTypeNode) => {
-        const typeName = extensionTypeNode.expression.getText(sourceFile);
+  const handleTypeNode = (typeNode: ts.Node) => {
+    if (ts.isParenthesizedTypeNode(typeNode)) {
+      typeNode = typeNode.type;
+    }
 
-        referenceTypeNames.add(typeName);
+    if (ts.isTypeReferenceNode(typeNode)) {
+      handleTypeReferenceNode(typeNode);
+    } else if (ts.isArrayTypeNode(typeNode)) {
+      handleTypeNode(typeNode.elementType);
+    } else if (ts.isTypeLiteralNode(typeNode)) {
+      typeNode.forEachChild(visitorExtract);
+    } else if (
+      ts.isIntersectionTypeNode(typeNode) ||
+      ts.isUnionTypeNode(typeNode)
+    ) {
+      typeNode.types.forEach((childNode: ts.TypeNode) => {
+        if (ts.isTypeReferenceNode(childNode)) {
+          handleTypeReferenceNode(childNode);
+        } else childNode.forEachChild(visitorExtract);
       });
-    });
-  }
+    }
+  };
 
-  const visitorExtract = (child: ts.Node) => {
-    const childNode = child as ts.PropertySignature;
-    if (!ts.isPropertySignature(childNode)) {
-      return;
+  const handleTypeReferenceNode = (typeRefNode: ts.TypeReferenceNode) => {
+    const typeName = typeRefNode.typeName.getText(sourceFile);
+    if (typeScriptHelper.indexOf(typeName) > -1 && typeRefNode.typeArguments) {
+      typeRefNode.typeArguments.forEach((t) => handleTypeNode(t));
+    } else {
+      referenceTypeNames.add(typeName);
     }
+  };
 
-    if (childNode.type) {
-      if (ts.isTypeReferenceNode(childNode.type)) {
-        referenceTypeNames.add(childNode.type.getText(sourceFile));
-      } else if (
-        ts.isArrayTypeNode(childNode.type) &&
-        ts.isTypeNode(childNode.type.elementType)
-      ) {
-        referenceTypeNames.add(childNode.type.elementType.getText(sourceFile));
-      } else if (ts.isTypeLiteralNode(childNode.type)) {
-        childNode.type.forEachChild(visitorExtract);
-      } else if (
-        ts.isIntersectionTypeNode(childNode.type) ||
-        ts.isUnionTypeNode(childNode.type)
-      ) {
-        childNode.type.types.forEach((typeNode: ts.TypeNode) => {
-          if (ts.isTypeReferenceNode(typeNode)) {
-            referenceTypeNames.add(typeNode.getText(sourceFile));
-          } else typeNode.forEachChild(visitorExtract);
+  if (ts.isInterfaceDeclaration(node)) {
+    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);
+
+          referenceTypeNames.add(typeName);
         });
-      }
+      });
     }
-  };
 
-  node.forEachChild(visitorExtract);
+    node.forEachChild(visitorExtract);
+  } else if (ts.isTypeAliasDeclaration(node)) {
+    handleTypeNode(node.type);
+  }
 
   return Array.from(referenceTypeNames);
 }