refactor: improve highlighted function (#4967)

packages/vitest/src/node/hoistMocks.ts

@@ -5,7 +5,8 @@ import type { AwaitExpression, CallExpression, Identifier, ImportDeclaration, Va
 import { findNodeAround } from 'acorn-walk'
 import type { PluginContext } from 'rollup'
 import { esmWalker } from '@vitest/utils/ast'
-import { highlight } from '@vitest/utils'
+import type { Colors } from '@vitest/utils'
+import { highlightCode } from '../utils/colors'
 import { generateCodeFrame } from './error'
 
 export type Positioned<T> = T & {
@@ -62,7 +63,7 @@ export function getBetterEnd(code: string, node: Node) {
 const regexpHoistable = /[ \t]*\b(vi|vitest)\s*\.\s*(mock|unmock|hoisted)\(/
 const hashbangRE = /^#!.*\n/
 
-export function hoistMocks(code: string, id: string, parse: PluginContext['parse']) {
+export function hoistMocks(code: string, id: string, parse: PluginContext['parse'], colors?: Colors) {
  const needHoisting = regexpHoistable.test(code)
 
  if (!needHoisting)
@@ -256,7 +257,7 @@ export function hoistMocks(code: string, id: string, parse: PluginContext['parse
  name: 'SyntaxError',
  message: _error.message,
  stack: _error.stack,
- frame: generateCodeFrame(highlight(code), 4, insideCall.start + 1),
+ frame: generateCodeFrame(highlightCode(id, code, colors), 4, insideCall.start + 1),
  }
  throw error
  }

packages/vitest/src/node/logger.ts

@@ -1,11 +1,10 @@
 import { createLogUpdate } from 'log-update'
 import c from 'picocolors'
-import { highlight } from '@vitest/utils'
-import { extname } from 'pathe'
 import { version } from '../../../../package.json'
 import type { ErrorWithDiff } from '../types'
 import type { TypeCheckError } from '../typecheck/typechecker'
 import { toArray } from '../utils'
+import { highlightCode } from '../utils/colors'
 import { divider } from './reporters/renderers/utils'
 import { RandomSequencer } from './sequencers/RandomSequencer'
 import type { Vitest } from './core'
@@ -23,14 +22,6 @@ const ERASE_DOWN = `${ESC}J`
 const ERASE_SCROLLBACK = `${ESC}3J`
 const CURSOR_TO_START = `${ESC}1;1H`
 const CLEAR_SCREEN = '\x1Bc'
-const HIGHLIGHT_SUPPORTED_EXTS = new Set(['js', 'ts'].flatMap(lang => [
- `.${lang}`,
- `.m${lang}`,
- `.c${lang}`,
- `.${lang}x`,
- `.m${lang}x`,
- `.c${lang}x`,
-]))
 
 export class Logger {
  outputStream = process.stdout
@@ -112,11 +103,7 @@ export class Logger {
  highlight(filename: string, source: string) {
  if (this._highlights.has(filename))
  return this._highlights.get(filename)!
- const ext = extname(filename)
- if (!HIGHLIGHT_SUPPORTED_EXTS.has(ext))
- return source
- const isJsx = ext.endsWith('x')
- const code = highlight(source, { jsx: isJsx, colors: c })
+ const code = highlightCode(filename, source)
  this._highlights.set(filename, code)
  return code
  }

packages/vitest/src/utils/colors.ts

@@ -0,0 +1,25 @@
+import type { Colors } from '@vitest/utils'
+import { highlight } from '@vitest/utils'
+import { extname } from 'pathe'
+import c from 'picocolors'
+
+const HIGHLIGHT_SUPPORTED_EXTS = new Set(['js', 'ts'].flatMap(lang => [
+ `.${lang}`,
+ `.m${lang}`,
+ `.c${lang}`,
+ `.${lang}x`,
+ `.m${lang}x`,
+ `.c${lang}x`,
+]))
+
+export function highlightCode(
+ id: string,
+ source: string,
+ colors?: Colors,
+) {
+ const ext = extname(id)
+ if (!HIGHLIGHT_SUPPORTED_EXTS.has(ext))
+ return source
+ const isJsx = ext.endsWith('x')
+ return highlight(source, { jsx: isJsx, colors: colors || c })
+}

test/core/test/injector-mock.test.ts

@@ -1,6 +1,7 @@
 import { parseAst } from 'rollup/parseAst'
 import { describe, expect, it, test } from 'vitest'
 import stripAnsi from 'strip-ansi'
+import { getDefaultColors } from '@vitest/utils'
 import { hoistMocks } from '../../../packages/vitest/src/node/hoistMocks'
 
 function parse(code: string, options: any) {
@@ -1205,7 +1206,7 @@ await vi
 describe('throws an error when nodes are incompatible', () => {
  const getErrorWhileHoisting = (code: string) => {
  try {
- hoistSimpleCode(code)
+ hoistMocks(code, '/test.js', parse, getDefaultColors())?.code.trim()
  }
  catch (err: any) {
  return err