feat: suppport ast transform hooks

Author: antfu

src/core/macros.ts

@@ -13,7 +13,7 @@ import {
 } from '@babel/types'
 import { types as t } from '@babel/core'
 import { parseExpression } from '@babel/parser'
-import { PropTypeData } from './types'
+import { PropTypeData } from '../types'
 
 // Special compiler macros
 const DEFINE_PROPS = 'defineProps'

src/core/parseSFC.ts

@@ -1,10 +1,10 @@
 import { Parser as HTMLParser } from 'htmlparser2'
 import { parse, ParserOptions } from '@babel/parser'
 import { camelize, capitalize, isHTMLTag, isSVGTag, isVoidTag } from '@vue/shared'
-import { ParsedSFC, ScriptTagMeta } from './types'
+import { ParsedSFC, ScriptSetupTransformOptions, ScriptTagMeta } from '../types'
 import { getIdentifierUsages } from './identifiers'
 
-export function parseSFC(code: string, id?: string): ParsedSFC {
+export function parseSFC(code: string, id?: string, options?: ScriptSetupTransformOptions): ParsedSFC {
   const components = new Set<string>()
   const expressions = new Set<string>()
   const identifiers = new Set<string>()
@@ -138,6 +138,9 @@ export function parseSFC(code: string, id?: string): ParsedSFC {
   scriptSetup.ast = parse(scriptSetup.content, parserOptions).program
   script.ast = parse(script.content || '', parserOptions).program
 
+  scriptSetup.ast = options?.astTransforms?.scriptSetup?.(scriptSetup.ast) || scriptSetup.ast
+  script.ast = options?.astTransforms?.script?.(script.ast) || script.ast
+
   return {
     id,
     template: {

src/core/transform.ts

@@ -1,11 +1,12 @@
 import MagicString from 'magic-string'
+import { ScriptSetupTransformOptions } from '../types'
 import { parseSFC } from './parseSFC'
 import { transformScriptSetup } from './transformScriptSetup'
 
-export function transform(input: string, id?: string) {
+export function transform(input: string, id?: string, options?: ScriptSetupTransformOptions) {
   const s = new MagicString(input)
   const sfc = parseSFC(input, id)
-  const { code } = transformScriptSetup(sfc)
+  const { code } = transformScriptSetup(sfc, options)
 
   const attributes = {
     ...sfc.script.attrs,

src/core/transformScriptSetup.ts

@@ -2,11 +2,11 @@ import { types as t } from '@babel/core'
 import { camelize, capitalize } from '@vue/shared'
 import traverse from '@babel/traverse'
 import generate from '@babel/generator'
-import { ParsedSFC } from './types'
+import { ParsedSFC, ScriptSetupTransformOptions } from '../types'
 import { applyMacros } from './macros'
 import { getIdentifierDeclarations } from './identifiers'
 
-export function transformScriptSetup(sfc: ParsedSFC) {
+export function transformScriptSetup(sfc: ParsedSFC, options?: ScriptSetupTransformOptions) {
   const { scriptSetup, script, template } = sfc
 
   const imports = scriptSetup.ast.body.filter(n => n.type === 'ImportDeclaration')
@@ -31,7 +31,7 @@ export function transformScriptSetup(sfc: ParsedSFC) {
     )
 
   // append `<script setup>` imports to `<script>`
-  const ast = t.program([
+  let ast = t.program([
     ...imports,
     ...script.ast.body,
   ])
@@ -139,7 +139,7 @@ export function transformScriptSetup(sfc: ParsedSFC) {
     )
   }
 
-  if (!hasBody) {
+  if (!hasBody && !options?.astTransforms) {
     return {
       ast: null,
       code: '',
@@ -152,6 +152,8 @@ export function transformScriptSetup(sfc: ParsedSFC) {
     t.exportDefaultDeclaration(__sfc) as any,
   )
 
+  ast = options?.astTransforms?.post?.(ast, sfc) || ast
+
   return {
     ast,
     code: generate(ast).code,

src/core/types.ts renamed to src/types.ts

@@ -28,9 +28,9 @@ export interface PropTypeData {
 }
 
 export interface ScriptSetupTransformOptions {
-
-}
-
-export interface ScriptSetupTransformContext {
-
+  astTransforms?: {
+    script?: (ast: Program) => Program
+    scriptSetup?: (ast: Program) => Program
+    post?: (ast: Program, sfc: ParsedSFC) => Program
+  }
 }

src/unplugin.ts

@@ -1,16 +1,17 @@
 import { createUnplugin } from 'unplugin'
+import { ScriptSetupTransformOptions } from './types'
 import { transform } from '.'
 
 const scriptSetupRE = /<script\s+setup/
 
-export default createUnplugin(() => ({
+export default createUnplugin<ScriptSetupTransformOptions>(options => ({
   name: 'vue2-script-setup-transform',
   enforce: 'pre',
   transformInclude(id) {
     return id.endsWith('.vue')
   },
-  transform(code) {
+  transform(code, id) {
     if (scriptSetupRE.test(code))
-      return transform(code)
+      return transform(code, id, options)
   },
 }))