Extremely rudimentary zod support.

package.json

@@ -11,6 +11,7 @@
     "test": "jest",
     "special": "ts-node src/index.ts '../../personal/Tangent-Mobile/api/entities/*.d.ts' -i -t -o ./schema/Tangent",
     "special2": "ts-node src/index.ts 'tests/comparisons/**/*.d.ts' -i -t -o ./schema/Tests/",
+    "zod": "ts-node src/index.ts 'tests/comparisons/single/zodStuff.ts' -i -t -z -o ./schema/Tests/ZodStuff.kt",
     "lint": "eslint . --ext .ts"
   },
   "author": "[email protected]",
@@ -31,15 +32,16 @@
     "nodemon": "^2.0.19",
     "prettier": "^2.4.1",
     "ts-jest": "^27.0.7",
-    "ts-node": "^10.4.0"
+    "ts-node": "^10.4.0",
+    "zod": "^3.23.4"
   },
   "dependencies": {
     "@typescript/vfs": "^1.4.0",
     "glob": "^7.2.0",
     "indent-string": "^4.0.0",
     "lodash": "^4.17.21",
     "promisify": "^0.0.3",
-    "yargs": "^17.2.1",
-    "typescript": "^4.4.4"
+    "typescript": "^4.4.4",
+    "yargs": "^17.2.1"
   }
 }

src/index.ts

@@ -49,6 +49,13 @@ const args = yargs
     default: false,
     describe: "Just throw in an import kotlinx.serialization.*",
   })
+  .option("experimentalZodSupport", {
+    alias: "z",
+    type: "boolean",
+    default: false,
+    describe:
+      "WIP support for Zod transpilations. Need to manually infer and expand",
+  })
   .showHelpOnFail(true)
   .help()
   .strict()
@@ -64,6 +71,7 @@ export type TranspileSingleArgs = {
   snakeToCamelCase: boolean;
   annotationNewLines: boolean;
   importStar: boolean;
+  experimentalZodSupport: boolean;
 };
 
 async function transpile(args: TranspileSingleArgs) {
@@ -74,6 +82,7 @@ async function transpile(args: TranspileSingleArgs) {
     snakeToCamelCase,
     annotationNewLines,
     importStar,
+    experimentalZodSupport,
   } = args;
   const config: MartokConfig = {
     files: resolve.files,
@@ -84,6 +93,7 @@ async function transpile(args: TranspileSingleArgs) {
       snakeToCamelCase,
       annotationNewLines,
       importStar,
+      experimentalZodSupport,
     },
   };
   const martok = new Martok(config);

src/martok/Martok.ts

@@ -20,6 +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/zod/ZodProcessor";
 
 type MartokState = {
   nameScope: string[];
@@ -37,6 +38,8 @@ export class Martok {
 
   public readonly imports;
 
+  public readonly zodProcessor: ZodProcessor;
+
   public get checker(): TypeChecker {
     return this.program.getTypeChecker();
   }
@@ -70,6 +73,10 @@ 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);
 
@@ -157,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: [
@@ -195,7 +204,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 {
+    } else if (!this.zodProcessor.allowImportThrough(file)) {
       throw new Error(
         `${file.fileName} is not within the given source root, it can't be included in this project.`
       );

src/martok/MartokOptions.ts

@@ -3,4 +3,5 @@ export type MartokOptions = {
   snakeToCamelCase?: boolean;
   annotationNewLines?: boolean;
   importStar?: boolean;
+  experimentalZodSupport?: boolean;
 };

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 type ${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);
+  }
+}

src/tests/comparisons.spec.ts

@@ -24,6 +24,7 @@ describe("Single File Comparisons", () => {
         sourceRoot: root,
         options: {
           dedupeTaggedUnions: true,
+          experimentalZodSupport: true,
         },
       });
       const out = sanitizeComparison(martok.generateMultiFile());

tests/comparisons/single/ZodStuff.kt

@@ -0,0 +1,21 @@
+/****************************************************
+** DO NOT EDIT THIS FILE BY HAND!                  **
+** This file was automatically generated by Martok **
+** More info at https://github.com/asarazan/martok **
+*****************************************************/
+package example
+
+import kotlinx.serialization.*
+import kotlinx.serialization.json.*
+
+/**
+ * @expand
+ */
+@Serializable
+data class FormData(
+  val firstName: String,
+  val lastName: String,
+  val email: String,
+  val phone: String? = null,
+  val url: String? = null
+)

tests/comparisons/single/zodStuff.ts

@@ -0,0 +1,9 @@
+import { z } from "zod";
+
+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(),
+});