progress check-in

src/martok/Martok.ts

@@ -20,7 +20,7 @@ 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";
+import { ZodProcessor } from "./processing/zod/ZodProcessor";
 
 type MartokState = {
   nameScope: string[];
@@ -73,9 +73,12 @@ export class Martok {
 
     // Create initial program
     this.program = this.compiler.compileFiles(fsMap);
+
+    this.zodProcessor = new ZodProcessor(this);
+    this.program = this.zodProcessor.modifyProgram();
+
     this.program = new TypeExpander(this).expand();
     this.imports = new ImportGenerator(this);
-    this.zodProcessor = new ZodProcessor(this);
 
     this.declarations = new DeclarationGenerator(this);
     this.storage = new AsyncLocalStorage<MartokState>();
@@ -161,12 +164,14 @@ export class Martok {
 
   private processFile(file: SourceFile): MartokOutFile {
     console.log(`Process File: ${file.fileName}...`);
+    // file = this.zodProcessor.processFile(file);
     const name = TsHelper.getBaseFileName(file.fileName);
     const pkg = this.getFilePackage(file);
     this.pushNameScope(pkg);
     const base: MartokOutFile = {
       name,
       pkg,
+      file,
       text: {
         package: `package ${pkg}`,
         imports: [

src/martok/MartokOutFile.ts

@@ -1,9 +1,11 @@
 import { kotlin } from "../kotlin/Klass";
 import Klass = kotlin.Klass;
+import ts from "typescript";
 
 export type MartokOutFile = {
   name: string;
   pkg: string;
+  file: ts.SourceFile;
   text: {
     package: string;
     imports: (string | null)[];

src/martok/processing/zod/MartokZodObject.ts

@@ -0,0 +1,7 @@
+export type MartokZodObject = {
+  identifier: string;
+  isExport: boolean;
+  pos: number;
+  end: number;
+  fullText: string;
+};

src/martok/processing/zod/ZodProcessor.ts

@@ -0,0 +1,80 @@
+import ts, {
+  isCallExpression,
+  isIdentifier,
+  isPropertyAccessExpression,
+  isVariableStatement,
+  SourceFile,
+  SyntaxKind,
+} from "typescript";
+import { Martok } from "../../Martok";
+import { MartokZodObject } from "./MartokZodObject";
+import _ from "lodash";
+
+export class ZodProcessor {
+  public constructor(private readonly martok: Martok) {}
+
+  public allowImportThrough(file: ts.SourceFile): boolean {
+    if (!this.martok.config.options?.experimentalZodSupport) return false;
+    return file.fileName.includes("/martok/node_modules/zod/lib/");
+  }
+
+  private zodObjects(file: ts.SourceFile): MartokZodObject[] {
+    const result: MartokZodObject[] = [];
+    for (const statement of file.statements) {
+      if (!isVariableStatement(statement)) continue;
+      const decl = statement.declarationList.declarations[0];
+      const initializer = decl.initializer;
+      if (!initializer) continue;
+      if (!isCallExpression(initializer)) continue;
+      const expression = initializer.expression;
+      if (!isPropertyAccessExpression(expression)) continue;
+      if (expression.expression.getText() !== "z") continue;
+      if (expression.name.getText() !== "object") continue;
+      const { pos, end } = statement;
+      const isExport = _.some(
+        statement.modifiers,
+        (value) => value.kind == SyntaxKind.ExportKeyword
+      );
+      const identifier = decl.name.getText();
+      const fullText = statement.getFullText();
+      result.push({
+        identifier,
+        isExport,
+        pos,
+        end,
+        fullText,
+      });
+    }
+    return result;
+  }
+
+  private stringReplace(zod: MartokZodObject): string {
+    const fullText = zod.fullText;
+    const renamed = fullText.replace(zod.identifier, `__${zod.identifier}`);
+    return `${renamed}
+/** 
+ * @expand 
+ **/ 
+export ${zod.identifier} = z.infer<typeof __${zod.identifier}>`;
+  }
+
+  private getText(file: SourceFile): string {
+    const zods = this.zodObjects(file);
+    if (!zods.length) return file.getFullText();
+    let fullText = file.getFullText();
+    for (const obj of _.reverse(zods)) {
+      fullText = fullText.replace(obj.fullText, this.stringReplace(obj));
+    }
+    console.log(fullText);
+    return fullText;
+  }
+
+  public modifyProgram(): ts.Program {
+    const fs = new Map<string, string>();
+    for (const fileName of this.martok.config.files) {
+      const sourceFile = this.martok.program.getSourceFile(fileName)!;
+      fs.set(fileName, this.getText(sourceFile));
+    }
+    return this.martok.compiler.compileFiles(fs);
+  }
+}

tests/comparisons/single/zodStuff.ts

@@ -7,8 +7,3 @@ export const FormData = z.object({
   email: z.string().email(),
   url: z.string().url().optional(),
 });
-
-/**
- * @expand
- */
-export type FormData = z.infer<typeof FormData>;