progress check-in

src/martok/Martok.ts

@@ -20,14 +20,12 @@ import { processSnakeCase } from "./processing/SnakeCase";
 import { processOldNames, sanitizeName } from "./processing/SanitizeNames";
 import { TypeExpander } from "./processing/TypeExpander";
 import { TsCompiler } from "./TsCompiler";
-import { ZodProcessor } from "./processing/ZodProcessor";
 
 type MartokState = {
   nameScope: string[];
   externalStatements: ts.Symbol[];
   additionalDeclarations: string[];
   typeReplacer: TypeReplacer;
-  zodProcessor: ZodProcessor;
 };
 
 export class Martok {
@@ -54,9 +52,6 @@ export class Martok {
   public get typeReplacer(): TypeReplacer {
     return this.storage.getStore()!.typeReplacer;
   }
-  public get zodProcessor(): ZodProcessor {
-    return this.storage.getStore()!.zodProcessor;
-  }
 
   private readonly storage;
   private readonly formatter;
@@ -121,7 +116,6 @@ export class Martok {
       externalStatements: [],
       additionalDeclarations: [],
       typeReplacer: new TypeReplacer(this),
-      zodProcessor: new ZodProcessor(this),
     };
     return this.storage.run(state, () => {
       const result = _(this.config.files)
@@ -179,10 +173,10 @@ export class Martok {
         declarations: [],
       },
     };
+    const statements = [...file.statements];
     base.text.declarations.push(
-      ...this.declarations.generateDeclarations(file)
+      ...this.declarations.generateDeclarations(statements)
     );
-    const statements = [...file.statements];
     const symbols = [...this.externalSymbols];
     const imports = _.uniq([
       ...this.imports.generateImports(statements),
@@ -201,7 +195,7 @@ export class Martok {
     let relativePath = path.resolve(path.dirname(file.fileName));
     if (relativePath.startsWith(this.config.sourceRoot)) {
       relativePath = relativePath.slice(this.config.sourceRoot.length);
-    } else if (!this.zodProcessor.isZodImport(file)) {
+    } else {
       throw new Error(
         `${file.fileName} is not within the given source root, it can't be included in this project.`
       );

src/martok/declarations/DeclarationGenerator.ts

@@ -11,9 +11,9 @@ export class DeclarationGenerator {
 
   public constructor(private readonly martok: Martok) {}
 
-  public generateDeclarations(file: SourceFile): (Klass | string)[] {
+  public generateDeclarations(statements: Statement[]): (Klass | string)[] {
     return _.compact(
-      file.statements.flatMap((value) => this.generateDeclaration(value))
+      statements.flatMap((value) => this.generateDeclaration(value))
     );
   }
 

src/martok/processing/TypeExpander.ts

@@ -2,6 +2,7 @@ import ts, { SourceFile, Statement, getJSDocTags } from "typescript";
 import { hasTypeArguments } from "../../typescript/utils";
 import { Martok } from "../Martok";
 import { extractJsDocs, insertJsDocs, JsDocProperty } from "./Comments";
+import { ZodProcessor } from "./ZodProcessor";
 
 type ExpandFile = {
   fileName: string;
@@ -15,7 +16,10 @@ type ExpandStatementWithImports = {
 };
 
 export class TypeExpander {
-  public constructor(private martok: Martok) {}
+  private zodProcessor: ZodProcessor;
+  public constructor(private martok: Martok) {
+    this.zodProcessor = new ZodProcessor(martok);
+  }
 
   private shouldIgnore(statement: Statement): boolean {
     return (
@@ -26,6 +30,7 @@ export class TypeExpander {
   }
 
   private shouldExpand(statement: Statement): boolean {
+    if (this.zodProcessor.shouldExpand(statement)) return true;
     if (!ts.isTypeAliasDeclaration(statement)) return false;
 
     const hasFlattenTag =
@@ -106,8 +111,11 @@ export class TypeExpander {
     if (this.shouldIgnore(statement)) return "";
     if (!this.shouldExpand(statement)) return statement.getFullText();
 
+    const type =
+      this.zodProcessor.getType(statement) ??
+      this.martok.checker.getTypeAtLocation(statement);
     const newExpandedType = this.martok.checker.typeToString(
-      this.martok.checker.getTypeAtLocation(statement),
+      type,
       statement,
       ts.TypeFormatFlags.InTypeAlias |
         ts.TypeFormatFlags.NoTypeReduction |

src/martok/processing/ZodProcessor.ts

@@ -1,14 +1,24 @@
-import ts, { isVariableStatement } from "typescript";
+import ts, { isVariableDeclaration, isVariableStatement } from "typescript";
 import { Martok } from "../Martok";
 
 export class ZodProcessor {
   public constructor(private martok: Martok) {}
 
-  public isZodImport(file: ts.SourceFile): boolean {
-    return file.fileName.includes("/martok/node_modules/zod/lib/");
+  public getType(statement: ts.Statement): ts.Type | undefined {
+    if (!isVariableStatement(statement)) return undefined;
+    try {
+      const decl = statement.declarationList.declarations[0];
+      if (!isVariableDeclaration(decl)) return undefined;
+      const type = this.martok.checker.getTypeAtLocation(decl);
+      const symbol = type.aliasSymbol ?? type.symbol;
+      return symbol.getEscapedName() === "ZodObject" ? type : undefined;
+    } catch (e: unknown) {
+      console.error(e);
+      return undefined;
+    }
   }
 
-  public isZodStatement(node: ts.Node): boolean {
-    return true; // TODO
+  public shouldExpand(statement: ts.Statement) {
+    return this.getType(statement) !== undefined;
   }
 }

tests/comparisons/single/zodStuff.ts

@@ -1,9 +1,19 @@
 import { z } from "zod";
 
-const FormData = z.object({
+export const FormData = z.object({
   firstName: z.string().min(1).max(18),
   lastName: z.string().min(1).max(18),
   phone: z.string().min(10).max(14).optional(),
   email: z.string().email(),
   url: z.string().url().optional(),
 });
+
+/**
+ * @expand
+ */
+export type FormDataType = typeof FormData.shape;
+
+/**
+ * @expand
+ */
+export type Inferred = z.infer<typeof FormData>;