fix(transform): use mlly to respect how users import lib

package.json

@@ -35,6 +35,7 @@
   "dependencies": {
     "estree-walker": "^3.0.1",
     "magic-string": "^0.26.2",
+    "mlly": "^0.5.4",
     "ufo": "^0.8.5",
     "unplugin": "^0.7.2"
   },

src/transform.ts

@@ -1,4 +1,4 @@
-import { createContext, runInContext } from 'node:vm'
+import { Context, createContext, runInContext } from 'node:vm'
 import { pathToFileURL } from 'node:url'
 
 import { walk } from 'estree-walker'
@@ -7,14 +7,14 @@ import type { SimpleCallExpression } from 'estree'
 import { createUnplugin } from 'unplugin'
 import MagicString from 'magic-string'
 import { parseURL, parseQuery } from 'ufo'
+import { findStaticImports, parseStaticImport } from 'mlly'
 
 import * as magicRegExp from 'magic-regexp'
 
 export const MagicRegExpTransformPlugin = createUnplugin(() => {
-  const context = createContext(magicRegExp)
-
   return {
     name: 'MagicRegExpTransformPlugin',
+    enforce: 'post',
     transformInclude(id) {
       const { pathname, search } = parseURL(decodeURIComponent(pathToFileURL(id).href))
       const { type } = parseQuery(search)
@@ -30,14 +30,49 @@ export const MagicRegExpTransformPlugin = createUnplugin(() => {
       }
     },
     transform(code, id) {
-      if (!code.includes('createRegExp')) return
+      if (!code.includes('magic-regexp')) return
+
+      const statements = findStaticImports(code).filter(i => i.specifier === 'magic-regexp')
+      if (!statements.length) return
+
+      const contextMap: Context = { ...magicRegExp }
+      const wrapperNames = []
+      let namespace: string
+
+      for (const i of statements.flatMap(i => parseStaticImport(i))) {
+        if (i.namespacedImport) {
+          namespace = i.namespacedImport
+          contextMap[i.namespacedImport] = magicRegExp
+        }
+        if (i.namedImports) {
+          for (const key in i.namedImports) {
+            contextMap[i.namedImports[key]] = magicRegExp[key]
+          }
+          if (i.namedImports.createRegExp) {
+            wrapperNames.push(i.namedImports.createRegExp)
+          }
+        }
+      }
+
+      const context = createContext(contextMap)
 
       const s = new MagicString(code)
 
       walk(this.parse(code), {
         enter(node: SimpleCallExpression) {
           if (node.type !== 'CallExpression') return
-          if ((node.callee as any).name !== 'createRegExp') return
+          if (
+            // Normal call
+            !wrapperNames.includes((node.callee as any).name) &&
+            // Namespaced call
+            (node.callee.type !== 'MemberExpression' ||
+              node.callee.object.type !== 'Identifier' ||
+              node.callee.object.name !== namespace ||
+              node.callee.property.type !== 'Identifier' ||
+              node.callee.property.name !== 'createRegExp')
+          ) {
+            return
+          }
 
           const { start, end } = node as any as { start: number; end: number }
 

test/transform.test.ts

@@ -5,32 +5,60 @@ import { MagicRegExpTransformPlugin } from '../src/transform'
 
 describe('transformer', () => {
   it('preserves context for dynamic regexps', () => {
-    expect(transform(`console.log(createRegExp(anyOf(keys)))`)).not.toBeDefined()
+    expect(
+      transform([
+        "import { createRegExp } from 'magic-regexp'",
+        `console.log(createRegExp(anyOf(keys)))`,
+      ])
+    ).not.toBeDefined()
   })
 
   it('statically replaces regexps where possible', () => {
     const code = transform([
+      "import { createRegExp, exactly, anyOf } from 'magic-regexp'",
+      '//', // this lets us tree-shake the import for use in our test-suite
       "const re1 = createRegExp(exactly('bar').notBefore('foo'))",
       "const re2 = createRegExp(anyOf(exactly('bar'), 'foo'))",
       "const re3 = createRegExp('/foo/bar')",
       // This line will be double-escaped in the snapshot
       "re3.test('/foo/bar')",
     ])
     expect(code).toMatchInlineSnapshot(`
-      "const re1 = /bar(?!foo)/
+      "import { createRegExp, exactly, anyOf } from 'magic-regexp'
+      //
+      const re1 = /bar(?!foo)/
       const re2 = /(bar|foo)/
       const re3 = /\\\\/foo\\\\/bar/
       re3.test('/foo/bar')"
     `)
     // ... but we test it here.
-    expect(eval(code)).toMatchInlineSnapshot('true')
+    expect(eval(code.split('//').pop())).toMatchInlineSnapshot('true')
+  })
+
+  it('respects how users import library', () => {
+    const code = transform([
+      "import { createRegExp as cRE } from 'magic-regexp'",
+      'import { exactly as ext, createRegExp } from "magic-regexp"',
+      'import * as magicRE from "magic-regexp"',
+      "const re1 = cRE(ext('bar').notBefore('foo'))",
+      "const re2 = magicRE.createRegExp(magicRE.anyOf('bar', 'foo'))",
+      "const re3 = createRegExp('test/value')",
+    ])
+    expect(code).toMatchInlineSnapshot(`
+      "import { createRegExp as cRE } from 'magic-regexp'
+      import { exactly as ext, createRegExp } from \\"magic-regexp\\"
+      import * as magicRE from \\"magic-regexp\\"
+      const re1 = /bar(?!foo)/
+      const re2 = /(bar|foo)/
+      const re3 = /test\\\\/value/"
+    `)
   })
 })
 
 const transform = (code: string | string[]) => {
   const plugin = MagicRegExpTransformPlugin.vite()
   return plugin.transform.call(
-    { parse: (code: string) => parse(code, { ecmaVersion: 2022 }) },
+    { parse: (code: string) => parse(code, { ecmaVersion: 2022, sourceType: 'module' }) },
     Array.isArray(code) ? code.join('\n') : code,
     'some-id.js'
   )?.code