|
| 1 | +import MagicString from 'magic-string' |
| 2 | +import { extract_names as extractNames } from 'periscopic' |
| 3 | +import type { Expression, ImportDeclaration } from 'estree' |
| 4 | +import type { AcornNode } from 'rollup' |
| 5 | +import type { Node, Positioned } from './esmWalker' |
| 6 | +import { esmWalker, isInDestructuringAssignment, isNodeInPattern, isStaticProperty } from './esmWalker' |
| 7 | + |
| 8 | +const viInjectedKey = '__vi_inject__' |
| 9 | +// const viImportMetaKey = '__vi_import_meta__' // to allow overwrite |
| 10 | +const viExportAllHelper = '__vi_export_all__' |
| 11 | + |
| 12 | +const skipHijack = [ |
| 13 | + '/@vite/client', |
| 14 | + '/@vite/env', |
| 15 | + /vite\/dist\/client/, |
| 16 | +] |
| 17 | + |
| 18 | +interface Options { |
| 19 | + cacheDir: string |
| 20 | +} |
| 21 | + |
| 22 | +// this is basically copypaste from Vite SSR |
| 23 | +// this method transforms all import and export statements into `__vi_injected__` variable |
| 24 | +// to allow spying on them. this can be disabled by setting `slowHijackESM` to `false` |
| 25 | +export function injectVitestModule(code: string, id: string, parse: (code: string, options: any) => AcornNode, options: Options) { |
| 26 | + if (skipHijack.some(skip => id.match(skip))) |
| 27 | + return |
| 28 | + |
| 29 | + const s = new MagicString(code) |
| 30 | + |
| 31 | + let ast: any |
| 32 | + try { |
| 33 | + ast = parse(code, { |
| 34 | + sourceType: 'module', |
| 35 | + ecmaVersion: 'latest', |
| 36 | + locations: true, |
| 37 | + }) |
| 38 | + } |
| 39 | + catch (err) { |
| 40 | + console.error(`Cannot parse ${id}:\n${(err as any).message}`) |
| 41 | + return |
| 42 | + } |
| 43 | + |
| 44 | + let uid = 0 |
| 45 | + const idToImportMap = new Map<string, string>() |
| 46 | + const declaredConst = new Set<string>() |
| 47 | + |
| 48 | + const hoistIndex = 0 |
| 49 | + |
| 50 | + let hasInjected = false |
| 51 | + |
| 52 | + // this will tranfrom import statements into dynamic ones, if there are imports |
| 53 | + // it will keep the import as is, if we don't need to mock anything |
| 54 | + // in browser environment it will wrap the module value with "vitest_wrap_module" function |
| 55 | + // that returns a proxy to the module so that named exports can be mocked |
| 56 | + const transformImportDeclaration = (node: ImportDeclaration) => { |
| 57 | + const source = node.source.value as string |
| 58 | + |
| 59 | + if (skipHijack.some(skip => source.match(skip))) |
| 60 | + return null |
| 61 | + |
| 62 | + const importId = `__vi_esm_${uid++}__` |
| 63 | + const hasSpecifiers = node.specifiers.length > 0 |
| 64 | + const code = hasSpecifiers |
| 65 | + ? `import { ${viInjectedKey} as ${importId} } from '${source}'\n` |
| 66 | + : `import '${source}'\n` |
| 67 | + return { |
| 68 | + code, |
| 69 | + id: importId, |
| 70 | + } |
| 71 | + } |
| 72 | + |
| 73 | + function defineImport(node: ImportDeclaration) { |
| 74 | + const declaration = transformImportDeclaration(node) |
| 75 | + if (!declaration) |
| 76 | + return null |
| 77 | + s.appendLeft(hoistIndex, declaration.code) |
| 78 | + return declaration.id |
| 79 | + } |
| 80 | + |
| 81 | + function defineImportAll(source: string) { |
| 82 | + const importId = `__vi_esm_${uid++}__` |
| 83 | + s.appendLeft(hoistIndex, `const { ${viInjectedKey}: ${importId} } = await import(${JSON.stringify(source)});\n`) |
| 84 | + return importId |
| 85 | + } |
| 86 | + |
| 87 | + function defineExport(position: number, name: string, local = name) { |
| 88 | + hasInjected = true |
| 89 | + s.appendLeft( |
| 90 | + position, |
| 91 | + `\nObject.defineProperty(${viInjectedKey}, "${name}", ` |
| 92 | + + `{ enumerable: true, configurable: true, get(){ return ${local} }});`, |
| 93 | + ) |
| 94 | + } |
| 95 | + |
| 96 | + // 1. check all import statements and record id -> importName map |
| 97 | + for (const node of ast.body as Node[]) { |
| 98 | + // import foo from 'foo' --> foo -> __import_foo__.default |
| 99 | + // import { baz } from 'foo' --> baz -> __import_foo__.baz |
| 100 | + // import * as ok from 'foo' --> ok -> __import_foo__ |
| 101 | + if (node.type === 'ImportDeclaration') { |
| 102 | + const importId = defineImport(node) |
| 103 | + if (!importId) |
| 104 | + continue |
| 105 | + s.remove(node.start, node.end) |
| 106 | + for (const spec of node.specifiers) { |
| 107 | + if (spec.type === 'ImportSpecifier') { |
| 108 | + idToImportMap.set( |
| 109 | + spec.local.name, |
| 110 | + `${importId}.${spec.imported.name}`, |
| 111 | + ) |
| 112 | + } |
| 113 | + else if (spec.type === 'ImportDefaultSpecifier') { |
| 114 | + idToImportMap.set(spec.local.name, `${importId}.default`) |
| 115 | + } |
| 116 | + else { |
| 117 | + // namespace specifier |
| 118 | + idToImportMap.set(spec.local.name, importId) |
| 119 | + } |
| 120 | + } |
| 121 | + } |
| 122 | + } |
| 123 | + |
| 124 | + // 2. check all export statements and define exports |
| 125 | + for (const node of ast.body as Node[]) { |
| 126 | + // named exports |
| 127 | + if (node.type === 'ExportNamedDeclaration') { |
| 128 | + if (node.declaration) { |
| 129 | + if ( |
| 130 | + node.declaration.type === 'FunctionDeclaration' |
| 131 | + || node.declaration.type === 'ClassDeclaration' |
| 132 | + ) { |
| 133 | + // export function foo() {} |
| 134 | + defineExport(node.end, node.declaration.id!.name) |
| 135 | + } |
| 136 | + else { |
| 137 | + // export const foo = 1, bar = 2 |
| 138 | + for (const declaration of node.declaration.declarations) { |
| 139 | + const names = extractNames(declaration.id as any) |
| 140 | + for (const name of names) |
| 141 | + defineExport(node.end, name) |
| 142 | + } |
| 143 | + } |
| 144 | + s.remove(node.start, (node.declaration as Node).start) |
| 145 | + } |
| 146 | + else { |
| 147 | + s.remove(node.start, node.end) |
| 148 | + if (node.source) { |
| 149 | + // export { foo, bar } from './foo' |
| 150 | + const importId = defineImportAll(node.source.value as string) |
| 151 | + // hoist re-exports near the defined import so they are immediately exported |
| 152 | + for (const spec of node.specifiers) { |
| 153 | + defineExport( |
| 154 | + hoistIndex, |
| 155 | + spec.exported.name, |
| 156 | + `${importId}.${spec.local.name}`, |
| 157 | + ) |
| 158 | + } |
| 159 | + } |
| 160 | + else { |
| 161 | + // export { foo, bar } |
| 162 | + for (const spec of node.specifiers) { |
| 163 | + const local = spec.local.name |
| 164 | + const binding = idToImportMap.get(local) |
| 165 | + defineExport(node.end, spec.exported.name, binding || local) |
| 166 | + } |
| 167 | + } |
| 168 | + } |
| 169 | + } |
| 170 | + |
| 171 | + // default export |
| 172 | + if (node.type === 'ExportDefaultDeclaration') { |
| 173 | + const expressionTypes = ['FunctionExpression', 'ClassExpression'] |
| 174 | + if ( |
| 175 | + 'id' in node.declaration |
| 176 | + && node.declaration.id |
| 177 | + && !expressionTypes.includes(node.declaration.type) |
| 178 | + ) { |
| 179 | + // named hoistable/class exports |
| 180 | + // export default function foo() {} |
| 181 | + // export default class A {} |
| 182 | + hasInjected = true |
| 183 | + const { name } = node.declaration.id |
| 184 | + s.remove(node.start, node.start + 15 /* 'export default '.length */) |
| 185 | + s.append( |
| 186 | + `\nObject.defineProperty(${viInjectedKey}, "default", ` |
| 187 | + + `{ enumerable: true, configurable: true, value: ${name} });`, |
| 188 | + ) |
| 189 | + } |
| 190 | + else { |
| 191 | + // anonymous default exports |
| 192 | + hasInjected = true |
| 193 | + s.update( |
| 194 | + node.start, |
| 195 | + node.start + 14 /* 'export default'.length */, |
| 196 | + `${viInjectedKey}.default =`, |
| 197 | + ) |
| 198 | + if (id.startsWith(options.cacheDir)) { |
| 199 | + // keep export default for optimized dependencies |
| 200 | + s.append(`\nexport default { ${viInjectedKey}: ${viInjectedKey}.default };\n`) |
| 201 | + } |
| 202 | + } |
| 203 | + } |
| 204 | + |
| 205 | + // export * from './foo' |
| 206 | + if (node.type === 'ExportAllDeclaration') { |
| 207 | + s.remove(node.start, node.end) |
| 208 | + const importId = defineImportAll(node.source.value as string) |
| 209 | + // hoist re-exports near the defined import so they are immediately exported |
| 210 | + if (node.exported) { |
| 211 | + defineExport(hoistIndex, node.exported.name, `${importId}`) |
| 212 | + } |
| 213 | + else { |
| 214 | + hasInjected = true |
| 215 | + s.appendLeft(hoistIndex, `${viExportAllHelper}(${viInjectedKey}, ${importId});\n`) |
| 216 | + } |
| 217 | + } |
| 218 | + } |
| 219 | + |
| 220 | + // 3. convert references to import bindings & import.meta references |
| 221 | + esmWalker(ast, { |
| 222 | + onIdentifier(id, parent, parentStack) { |
| 223 | + const grandparent = parentStack[1] |
| 224 | + const binding = idToImportMap.get(id.name) |
| 225 | + if (!binding) |
| 226 | + return |
| 227 | + |
| 228 | + if (isStaticProperty(parent) && parent.shorthand) { |
| 229 | + // let binding used in a property shorthand |
| 230 | + // { foo } -> { foo: __import_x__.foo } |
| 231 | + // skip for destructuring patterns |
| 232 | + if ( |
| 233 | + !isNodeInPattern(parent) |
| 234 | + || isInDestructuringAssignment(parent, parentStack) |
| 235 | + ) |
| 236 | + s.appendLeft(id.end, `: ${binding}`) |
| 237 | + } |
| 238 | + else if ( |
| 239 | + (parent.type === 'PropertyDefinition' |
| 240 | + && grandparent?.type === 'ClassBody') |
| 241 | + || (parent.type === 'ClassDeclaration' && id === parent.superClass) |
| 242 | + ) { |
| 243 | + if (!declaredConst.has(id.name)) { |
| 244 | + declaredConst.add(id.name) |
| 245 | + // locate the top-most node containing the class declaration |
| 246 | + const topNode = parentStack[parentStack.length - 2] |
| 247 | + s.prependRight(topNode.start, `const ${id.name} = ${binding};\n`) |
| 248 | + } |
| 249 | + } |
| 250 | + else { |
| 251 | + s.update(id.start, id.end, binding) |
| 252 | + } |
| 253 | + }, |
| 254 | + // TODO: make env updatable |
| 255 | + onImportMeta() { |
| 256 | + // s.update(node.start, node.end, viImportMetaKey) |
| 257 | + }, |
| 258 | + onDynamicImport(node) { |
| 259 | + const replace = '__vi_wrap_module__(import(' |
| 260 | + s.overwrite(node.start, (node.source as Positioned<Expression>).start, replace) |
| 261 | + s.overwrite(node.end - 1, node.end, '))') |
| 262 | + }, |
| 263 | + }) |
| 264 | + |
| 265 | + if (hasInjected) { |
| 266 | + // make sure "__vi_injected__" is declared as soon as possible |
| 267 | + s.prepend(`const ${viInjectedKey} = { [Symbol.toStringTag]: "Module" };\n`) |
| 268 | + s.append(`\nexport { ${viInjectedKey} }`) |
| 269 | + } |
| 270 | + |
| 271 | + return { |
| 272 | + ast, |
| 273 | + code: s.toString(), |
| 274 | + map: s.generateMap({ hires: true, source: id }), |
| 275 | + } |
| 276 | +} |
0 commit comments