From f24887d476043193f1ab9ef66de7261d01003b31 Mon Sep 17 00:00:00 2001 From: Aaron Sarazan Date: Thu, 2 May 2024 16:12:50 -0700 Subject: [PATCH] progress check-in --- src/martok/Martok.ts | 9 ++- src/martok/MartokOutFile.ts | 2 + src/martok/processing/ZodProcessor.ts | 10 --- src/martok/processing/zod/MartokZodObject.ts | 7 ++ src/martok/processing/zod/ZodProcessor.ts | 80 ++++++++++++++++++++ tests/comparisons/single/zodStuff.ts | 5 -- 6 files changed, 96 insertions(+), 17 deletions(-) delete mode 100644 src/martok/processing/ZodProcessor.ts create mode 100644 src/martok/processing/zod/MartokZodObject.ts create mode 100644 src/martok/processing/zod/ZodProcessor.ts diff --git a/src/martok/Martok.ts b/src/martok/Martok.ts index 106b4e8..04426dc 100644 --- a/src/martok/Martok.ts +++ b/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(); @@ -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: [ diff --git a/src/martok/MartokOutFile.ts b/src/martok/MartokOutFile.ts index 9d3bb83..7367446 100644 --- a/src/martok/MartokOutFile.ts +++ b/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)[]; diff --git a/src/martok/processing/ZodProcessor.ts b/src/martok/processing/ZodProcessor.ts deleted file mode 100644 index b638cec..0000000 --- a/src/martok/processing/ZodProcessor.ts +++ /dev/null @@ -1,10 +0,0 @@ -import ts from "typescript"; -import { Martok } from "../Martok"; - -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/"); - } -} diff --git a/src/martok/processing/zod/MartokZodObject.ts b/src/martok/processing/zod/MartokZodObject.ts new file mode 100644 index 0000000..4e98e0c --- /dev/null +++ b/src/martok/processing/zod/MartokZodObject.ts @@ -0,0 +1,7 @@ +export type MartokZodObject = { + identifier: string; + isExport: boolean; + pos: number; + end: number; + fullText: string; +}; diff --git a/src/martok/processing/zod/ZodProcessor.ts b/src/martok/processing/zod/ZodProcessor.ts new file mode 100644 index 0000000..ac772f6 --- /dev/null +++ b/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`; + } + + 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(); + 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); + } +} diff --git a/tests/comparisons/single/zodStuff.ts b/tests/comparisons/single/zodStuff.ts index ea507c5..ee618f7 100644 --- a/tests/comparisons/single/zodStuff.ts +++ b/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;