From 433a211b2a257038e54737f4d16e71c2cb62fded Mon Sep 17 00:00:00 2001 From: Aaron Sarazan Date: Mon, 29 Apr 2024 15:15:40 -0700 Subject: [PATCH] progress check-in --- src/martok/Martok.ts | 12 +++-------- .../declarations/DeclarationGenerator.ts | 4 ++-- src/martok/processing/TypeExpander.ts | 12 +++++++++-- src/martok/processing/ZodProcessor.ts | 20 ++++++++++++++----- tests/comparisons/single/zodStuff.ts | 12 ++++++++++- 5 files changed, 41 insertions(+), 19 deletions(-) diff --git a/src/martok/Martok.ts b/src/martok/Martok.ts index 0125ef5..d903924 100644 --- a/src/martok/Martok.ts +++ b/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.` ); diff --git a/src/martok/declarations/DeclarationGenerator.ts b/src/martok/declarations/DeclarationGenerator.ts index 83a363b..ebffc57 100644 --- a/src/martok/declarations/DeclarationGenerator.ts +++ b/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)) ); } diff --git a/src/martok/processing/TypeExpander.ts b/src/martok/processing/TypeExpander.ts index 97f1b76..216573a 100644 --- a/src/martok/processing/TypeExpander.ts +++ b/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 | diff --git a/src/martok/processing/ZodProcessor.ts b/src/martok/processing/ZodProcessor.ts index fb30a2b..529d250 100644 --- a/src/martok/processing/ZodProcessor.ts +++ b/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; } } diff --git a/tests/comparisons/single/zodStuff.ts b/tests/comparisons/single/zodStuff.ts index c83f50b..1daf537 100644 --- a/tests/comparisons/single/zodStuff.ts +++ b/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;