refactor: improve highlighted function (#4967)

Author: sheremet-va

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