genzod.ts

import fs from 'node:fs/promises'
import path from 'node:path'
import rimraf from 'rimraf'
import { glob } from 'glob'
import { generate } from 'ts-to-zod'
import ts from 'typescript'
import { v4 as uuidv4 } from 'uuid'

const OUTDIR = 'packages/types/gen'
const { factory } = ts

const processSource = (src: string) => {
  const sourceFile = ts.createSourceFile('temp.ts', src, ts.ScriptTarget.Latest)
  const importedVariables = new Map<string, string>()

  ts.forEachChild(sourceFile, node => {
    if (ts.isImportDeclaration(node) && node.importClause?.namedBindings) {
      const bindings = node.importClause.namedBindings
      if (ts.isNamedImports(bindings)) {
        bindings.elements.forEach(el => {
          if (!importedVariables.has(el.name.text)) {
            importedVariables.set(el.name.text, uuidv4())
          }
        })
      }
    }
  })

  const placeholderTransformer: ts.TransformerFactory<
    ts.SourceFile
  > = context => {
    const visit: ts.Visitor = node => {
      if (ts.isIdentifier(node) && importedVariables.has(node.text)) {
        return factory.createStringLiteral(
          importedVariables.get(node.text) as string
        )
      }
      return ts.visitEachChild(node, visit, context)
    }
    return node => ts.visitNode(node, visit)
  }

  const result = ts.transform(sourceFile, [placeholderTransformer])
  const printer = ts.createPrinter()
  const output = printer.printNode(
    ts.EmitHint.Unspecified,
    result.transformed[0],
    sourceFile
  )

  const literals = new Map<string, string>()
  for (const [key, value] of importedVariables.entries()) {
    literals.set(value, key)
  }

  return { output, literals }
}

const restoreModules = (src: string, literals) => {
  const restoreTransformer: ts.TransformerFactory<ts.SourceFile> = context => {
    const visit: ts.Visitor = node => {
      if (ts.isImportDeclaration(node)) {
        if (!ts.isStringLiteral(node.moduleSpecifier) || node.moduleSpecifier.text === 'zod') {
          return node
        }
        if (!node.importClause?.namedBindings || !ts.isNamedImports(node.importClause.namedBindings)) {
          throw new Error('Unsupported import declaration')
        }
        const namedImports = factory.createNamedImports(
          node.importClause.namedBindings.elements.map(el => {
            const base = el.name.text
            const schemaName = factory.createIdentifier(`${base}Schema`)
            const schemaSpec = factory.createImportSpecifier(
              false,
              undefined,
              schemaName
            )
            return schemaSpec
          })
        )
        const clause = factory.createImportClause(false, undefined, namedImports)
        const schemaImport = factory.createImportDeclaration(
          node.modifiers,
          clause,
          node.moduleSpecifier,
          undefined
        )
        return [node, schemaImport]
      }

      if (
        ts.isCallExpression(node) &&
        node.arguments.length === 1 &&
        ts.isStringLiteral(node.arguments[0]) &&
        literals.has(node.arguments[0].text)
      ) {
        const id = literals.get(node.arguments[0].text) as string
        return factory.createIdentifier(`${id}Schema`)
      }
      return ts.visitEachChild(node, visit, context)
    }

    return node => ts.visitNode(node, visit)
  }

  const sourceFile = ts.createSourceFile('temp.ts', src, ts.ScriptTarget.Latest)
  const result = ts.transform(sourceFile, [restoreTransformer])
  const printer = ts.createPrinter()
  return printer.printNode(
    ts.EmitHint.Unspecified,
    result.transformed[0],
    sourceFile
  )
}

const write = async (file: string, contents: string) => {
  const dir = path.dirname(file)
  await fs.mkdir(dir, { recursive: true })
  await fs.writeFile(file, contents)
}

const main = async () => {
  await rimraf(OUTDIR)

  const files = await glob('packages/types/src/**/*.ts')
  for (const file of files) {
    console.log('PROCESSING', file)

    const outfile = file.replace('types/src', 'types/gen')
    const raw = await fs.readFile(file, 'utf8')

    if (path.basename(file) === 'index.ts') {
      await write(outfile, raw)
      continue
    }

    const { output: src, literals } = processSource(raw)

    const gen = generate({
      sourceText: src,
      getSchemaName: id => `${id}Schema`,
    })

    // generate zod schemas
    let zods = gen.getZodSchemasFile('null')
    zods = zods.replace(/^import .*? "null";$/gm, '')

    // concat and write
    let build = [raw, zods].join('\n\n')
    build = restoreModules(build, literals)
    await write(outfile, build)
  }
}

main()