refactor!: improve API

package.json

@@ -47,9 +47,9 @@
       "require": "./dist/esbuild.js",
       "import": "./dist/esbuild.mjs"
     },
-    "./cores": {
-      "require": "./dist/cores.js",
-      "import": "./dist/cores.mjs"
+    "./presets": {
+      "require": "./dist/presets.js",
+      "import": "./dist/presets.mjs"
     },
     "./*": "./*"
   },
@@ -89,6 +89,7 @@
     "estree-walker": "^2.0.2",
     "fast-glob": "^3.2.11",
     "prettier": "^2.7.1",
+    "rollup": "^2.75.7",
     "tsup": "^6.1.2",
     "tsx": "^3.4.3",
     "typescript": "^4.7.4",

src/core/options.ts

@@ -1,8 +1,7 @@
 import { toArray } from '@antfu/utils'
-import type MagicString from 'magic-string'
+import type { Transformer } from './types'
 import type { ParserOptions } from '@babel/parser'
-import type { Arrayable, Awaitable } from '@antfu/utils'
-import type { Node } from '@babel/types'
+import type { Arrayable } from '@antfu/utils'
 import type { FilterPattern } from '@rollup/pluginutils'
 
 export interface Options {
@@ -13,20 +12,6 @@ export interface Options {
   transformer?: Arrayable<Transformer<any>>
 }
 
-export interface Transformer<T extends Node = Node> {
-  filterFile?: (id: string) => Awaitable<boolean>
-  filterNode?:
-    | ((node: Node, parent: Node) => Awaitable<boolean>)
-    | ((node: Node, parent: Node) => node is T)
-  transform: (
-    node: T,
-    code: MagicString,
-    context: {
-      id: string
-    }
-  ) => Awaitable<string | undefined | null | void>
-}
-
 export type OptionsResolved = Omit<Required<Options>, 'transformer'> & {
   transformer: Transformer<any>[]
 }

src/core/presets/remove-wrapper-function.ts

@@ -0,0 +1,22 @@
+import { toArray } from '@antfu/utils'
+import type { Arrayable } from '@antfu/utils'
+import type { CallExpression } from '@babel/types'
+import type { Transformer } from '../types'
+
+/**
+ * Removes wrapper function. e.g `defineComponent`, `defineConfig`...
+ * @param functionNames - function names to remove
+ * @returns Transformer
+ */
+export const RemoveWrapperFunction = (
+  functionNames: Arrayable<string>
+): Transformer<CallExpression> => ({
+  onNode: (node) =>
+    node.type === 'CallExpression' &&
+    node.callee.type === 'Identifier' &&
+    toArray(functionNames).includes(node.callee.name),
+
+  transform(node) {
+    return node.arguments[0]
+  },
+})

src/core/transform.ts

@@ -1,33 +1,46 @@
 import { walk } from 'estree-walker'
 import MagicString from 'magic-string'
 import { parseCode } from './parse'
-import type { TransformResult } from 'unplugin'
-import type { Node } from '@babel/types'
-import type { OptionsResolved, Transformer } from './options'
+import type { NodeRef, Transformer, TransformerParsed } from './types'
+import type { SourceMap } from 'rollup'
+import type { ExpressionStatement, Node } from '@babel/types'
+import type { OptionsResolved } from './options'
 
-export const transform = async (
-  code: string,
-  id: string,
-  options: Pick<OptionsResolved, 'parserOptions' | 'transformer'>
-): Promise<TransformResult> => {
-  interface TransformerParsed {
-    transformer: Transformer
-    nodes: Node[]
+const nodeRefs: Map<Node, NodeRef> = new Map()
+function getNodeRef(node: Node) {
+  if (nodeRefs.has(node)) return nodeRefs.get(node)!
+  const ref: NodeRef = {
+    value: node,
+    set(node: Node) {
+      this.value = node
+    },
   }
+  nodeRefs.set(node, ref)
+  return ref
+}
 
+async function getTransformersByFile(transformer: Transformer[], id: string) {
   const transformers = (
     await Promise.all(
-      options.transformer.map(
-        async (t): Promise<TransformerParsed | undefined> => {
-          if (t.filterFile && !(await t.filterFile(id))) return undefined
-          return {
-            transformer: t,
-            nodes: [],
-          }
+      transformer.map(async (t): Promise<TransformerParsed | undefined> => {
+        if (t.transformInclude && !(await t.transformInclude(id)))
+          return undefined
+        return {
+          transformer: t,
+          nodes: [],
         }
-      )
+      })
     )
   ).filter((t): t is TransformerParsed => !!t)
+  return transformers
+}
+
+export const transform = async (
+  code: string,
+  id: string,
+  options: Pick<OptionsResolved, 'parserOptions' | 'transformer'>
+): Promise<{ code: string; map: SourceMap } | undefined> => {
+  const transformers = await getTransformersByFile(options.transformer, id)
   if (transformers.length === 0) return
 
   const nodes = parseCode(code, id, options.parserOptions)
@@ -38,11 +51,11 @@ export const transform = async (
       if (!node.type) return
       const p = (async () => {
         for (const { transformer, nodes } of transformers) {
-          if (transformer.filterNode) {
-            const bool = await transformer.filterNode?.(node, parent)
+          if (transformer.onNode) {
+            const bool = await transformer.onNode?.(node, parent)
             if (!bool) continue
           }
-          nodes.push(node)
+          nodes.push(getNodeRef(node))
         }
       })()
       promises.push(p)
@@ -54,7 +67,32 @@ export const transform = async (
   const s = new MagicString(code)
   for (const { transformer, nodes } of transformers) {
     for (const node of nodes) {
-      await transformer.transform(node, s, { id })
+      const value = node.value
+      const result = await transformer.transform(value, code, { id })
+      if (!result) continue
+
+      let newAST: Node
+      if (typeof result === 'string') {
+        s.overwrite(value.start!, value.end!, result)
+        newAST = (
+          parseCode(
+            `(${result})`,
+            id,
+            options.parserOptions
+          )[0] as ExpressionStatement
+        ).expression
+        newAST.start = value.start!
+        newAST.end = value.end!
+      } else {
+        s.overwrite(
+          value.start!,
+          value.end!,
+          s.slice(result.start!, result.end!)
+        )
+        newAST = result
+      }
+
+      node.set(newAST)
     }
   }
   if (!s.hasChanged()) return undefined

src/core/types.ts

@@ -0,0 +1,26 @@
+import type { Awaitable } from '@antfu/utils'
+import type { Node } from '@babel/types'
+
+export interface TransformerParsed {
+  transformer: Transformer
+  nodes: NodeRef[]
+}
+
+export interface NodeRef {
+  value: Node
+  set: (node: Node) => void
+}
+
+export interface Transformer<T extends Node = Node> {
+  transformInclude?: (id: string) => Awaitable<boolean>
+  onNode?:
+    | ((node: Node, parent: Node) => Awaitable<boolean>)
+    | ((node: Node, parent: Node) => node is T)
+  transform: (
+    node: T,
+    code: string,
+    context: {
+      id: string
+    }
+  ) => Awaitable<string | Node | undefined | null | void>
+}

src/index.ts

@@ -26,3 +26,4 @@ export default createUnplugin<Options>((options = {}) => {
 export * from './core/options'
 export * from './core/parse'
 export * from './core/transform'
+export * from './core/types'

src/presets.ts

@@ -0,0 +1 @@
+export * from './core/presets/remove-wrapper-function'

tests/basic.test.ts

@@ -1,68 +1,78 @@
 import { expect, test } from 'vitest'
 import { transform } from '../src/core/transform'
-import type { CallExpression, Identifier, StringLiteral } from '@babel/types'
+import type { Identifier, NumericLiteral, StringLiteral } from '@babel/types'
 import type { OptionsResolved, Transformer } from '../src/core/options'
 
+const changeString: Transformer<StringLiteral> = {
+  onNode: (node): node is StringLiteral => node.type === 'StringLiteral',
+  transform() {
+    return "'Hello'"
+  },
+}
+
+const changeVarName: Transformer<Identifier> = {
+  onNode: (node): node is Identifier =>
+    node.type === 'Identifier' && node.name === 'foo',
+  transform() {
+    return 'newName'
+  },
+}
+
+const overwriteVarName: Transformer<Identifier> = {
+  onNode: (node): node is Identifier => node.type === 'Identifier',
+  transform(node) {
+    return `overwrite_${node.name}`
+  },
+}
+
+const timesTen: Transformer<NumericLiteral> = {
+  onNode: (node): node is NumericLiteral => node.type === 'NumericLiteral',
+  transform(node) {
+    return `${node.value * 10}`
+  },
+}
+
 test('basic', async () => {
-  const source = `const foo = 'string'`
+  const source = `const foo = 'string'\nlet i = 10`
   const options: Pick<OptionsResolved, 'parserOptions' | 'transformer'> = {
     transformer: [],
     parserOptions: {},
   }
-  let code = await transform(source, 'foo.js', options)
+  let code = (await transform(source, 'foo.js', options))?.code
   expect(code).toMatchInlineSnapshot('undefined')
 
-  const transformer: Transformer<StringLiteral> = {
-    filterNode: (node): node is StringLiteral => node.type === 'StringLiteral',
-    transform(node, code) {
-      code.overwrite(node.start!, node.end!, '"Hello"')
-    },
-  }
-  options.transformer.push(transformer)
-  code = await transform(source, 'foo.js', options)
-  expect(code).toMatchInlineSnapshot('"const foo = \\"Hello\\""')
+  options.transformer = [changeString]
+  code = (await transform(source, 'foo.js', options))?.code
+  expect(code).toMatchInlineSnapshot(`
+    "const foo = 'Hello'
+    let i = 10"
+  `)
 
-  const transformer2: Transformer<Identifier> = {
-    filterNode: (node): node is Identifier => node.type === 'Identifier',
-    transform(node, code) {
-      code.overwrite(node.start!, node.end!, 'newName')
-    },
-  }
-  options.transformer.push(transformer2)
-  code = await transform(source, 'foo.js', options)
-  expect(code).toMatchInlineSnapshot('"const newName = \\"Hello\\""')
+  options.transformer = [changeVarName]
+  code = (await transform(source, 'foo.js', options))?.code
+  expect(code).toMatchInlineSnapshot(`
+    "const newName = 'string'
+    let i = 10"
+  `)
 
-  options.transformer.splice(0, 1)
-  code = await transform(source, 'foo.js', options)
-  expect(code).toMatchInlineSnapshot('"const newName = \'string\'"')
-})
+  options.transformer = [changeString, changeVarName]
+  code = (await transform(source, 'foo.js', options))?.code
+  expect(code).toMatchInlineSnapshot(`
+    "const newName = 'Hello'
+    let i = 10"
+  `)
 
-test('remove wrapper function', async () => {
-  const source = `const comp = defineComponent({
-    render() {
-      return []
-    }
-  })`
-  const transformer: Transformer<CallExpression> = {
-    filterNode: (node) =>
-      node.type === 'CallExpression' &&
-      node.callee.type === 'Identifier' &&
-      node.callee.name === 'defineComponent',
-    transform(node, code) {
-      const [arg] = node.arguments
-      code.overwrite(node.start!, node.end!, code.slice(arg.start!, arg.end!))
-    },
-  }
-  const options: Pick<OptionsResolved, 'parserOptions' | 'transformer'> = {
-    transformer: [transformer],
-    parserOptions: {},
-  }
-  const code = await transform(source, 'foo.js', options)
+  options.transformer = [changeString, changeVarName, overwriteVarName]
+  code = (await transform(source, 'foo.js', options))?.code
+  expect(code).toMatchInlineSnapshot(`
+    "const overwrite_newName = 'Hello'
+    let overwrite_i = 10"
+  `)
+
+  options.transformer = [timesTen, timesTen, timesTen]
+  code = (await transform(source, 'foo.js', options))?.code
   expect(code).toMatchInlineSnapshot(`
-    "const comp = {
-        render() {
-          return []
-        }
-      }"
+    "const foo = 'string'
+    let i = 10000"
   `)
 })