feat: add syntax highlighting to error messages (#4813)

examples/mocks/package.json

@@ -18,11 +18,11 @@
     "tinyspy": "^0.3.2"
   },
   "devDependencies": {
-    "@vitest/ui": "latest",
+    "@vitest/ui": "workspace:*",
     "react": "^18.0.0",
     "sweetalert2": "^11.6.16",
     "vite": "latest",
-    "vitest": "latest",
+    "vitest": "workspace:*",
     "vue": "^3.3.8",
     "zustand": "^4.1.1"
   },

packages/utils/package.json

@@ -69,6 +69,7 @@
   },
   "devDependencies": {
     "@jridgewell/trace-mapping": "^0.3.20",
-    "@types/estree": "^1.0.5"
+    "@types/estree": "^1.0.5",
+    "tinyhighlight": "^0.3.2"
   }
 }

packages/utils/rollup.config.js

@@ -4,6 +4,7 @@ import esbuild from 'rollup-plugin-esbuild'
 import dts from 'rollup-plugin-dts'
 import resolve from '@rollup/plugin-node-resolve'
 import json from '@rollup/plugin-json'
+import commonjs from '@rollup/plugin-commonjs'
 
 const require = createRequire(import.meta.url)
 const pkg = require('./package.json')
@@ -32,6 +33,7 @@ const plugins = [
   esbuild({
     target: 'node14',
   }),
+  commonjs(),
 ]
 
 export default defineConfig([

packages/utils/src/colors.ts

@@ -27,16 +27,17 @@ const colorsMap = {
   bgWhite: ['\x1B[47m', '\x1B[49m'],
 } as const
 
-type ColorName = keyof typeof colorsMap
-type ColorsMethods = {
-  [Key in ColorName]: {
-    (input: unknown): string
-    open: string
-    close: string
-  }
+export type ColorName = keyof typeof colorsMap
+export interface ColorMethod {
+  (input: unknown): string
+  open: string
+  close: string
+}
+export type ColorsMethods = {
+  [Key in ColorName]: ColorMethod
 }
 
-type Colors = ColorsMethods & {
+export type Colors = ColorsMethods & {
   isColorSupported: boolean
   reset: (input: unknown) => string
 }

packages/utils/src/highlight.ts

@@ -0,0 +1,43 @@
+import { type TokenColors, highlight as baseHighlight } from 'tinyhighlight'
+import type { ColorName } from './colors'
+import { getColors } from './colors'
+
+type Colors = Record<ColorName, (input: string) => string>
+
+function getDefs(c: Colors): TokenColors {
+  const Invalid = (text: string) => c.white(c.bgRed(c.bold(text)))
+  return {
+    Keyword: c.magenta,
+    IdentifierCapitalized: c.yellow,
+    Punctuator: c.yellow,
+    StringLiteral: c.green,
+    NoSubstitutionTemplate: c.green,
+    MultiLineComment: c.gray,
+    SingleLineComment: c.gray,
+    RegularExpressionLiteral: c.cyan,
+    NumericLiteral: c.blue,
+    TemplateHead: text => c.green(text.slice(0, text.length - 2)) + c.cyan(text.slice(-2)),
+    TemplateTail: text => c.cyan(text.slice(0, 1)) + c.green(text.slice(1)),
+    TemplateMiddle: text => c.cyan(text.slice(0, 1)) + c.green(text.slice(1, text.length - 2)) + c.cyan(text.slice(-2)),
+    IdentifierCallable: c.blue,
+    PrivateIdentifierCallable: text => `#${c.blue(text.slice(1))}`,
+    Invalid,
+
+    JSXString: c.green,
+    JSXIdentifier: c.yellow,
+    JSXInvalid: Invalid,
+    JSXPunctuator: c.yellow,
+  }
+}
+
+interface HighlightOptions {
+  jsx?: boolean
+  colors?: Colors
+}
+
+export function highlight(code: string, options: HighlightOptions = { jsx: false }) {
+  return baseHighlight(code, {
+    jsx: options.jsx,
+    colors: getDefs(options.colors || getColors()),
+  })
+}

packages/utils/src/index.ts

@@ -8,3 +8,4 @@ export * from './constants'
 export * from './colors'
 export * from './base'
 export * from './offset'
+export * from './highlight'

packages/utils/src/types.ts

@@ -44,4 +44,5 @@ export interface ErrorWithDiff extends Error {
   type?: string
   frame?: string
   diff?: string
+  codeFrame?: string
 }

packages/vitest/src/node/core.ts

@@ -670,13 +670,15 @@ export class Vitest {
 
     const onChange = (id: string) => {
       id = slash(id)
+      this.logger.clearHighlightCache(id)
       updateLastChanged(id)
       const needsRerun = this.handleFileChanged(id)
       if (needsRerun.length)
         this.scheduleRerun(needsRerun)
     }
     const onUnlink = (id: string) => {
       id = slash(id)
+      this.logger.clearHighlightCache(id)
       this.invalidates.add(id)
 
       if (this.state.filesMap.has(id)) {

packages/vitest/src/node/error.ts

@@ -67,6 +67,8 @@ export async function printError(error: unknown, project: WorkspaceProject | und
   if (type)
     printErrorType(type, project.ctx)
   printErrorMessage(e, logger)
+  if (e.codeFrame)
+    logger.error(`${e.codeFrame}\n`)
 
   // E.g. AssertionError from assert does not set showDiff but has both actual and expected properties
   if (e.diff)
@@ -80,7 +82,7 @@ export async function printError(error: unknown, project: WorkspaceProject | und
     printStack(project, stacks, nearest, errorProperties, (s) => {
       if (showCodeFrame && s === nearest && nearest) {
         const sourceCode = readFileSync(nearest.file, 'utf-8')
-        logger.error(generateCodeFrame(sourceCode, 4, s))
+        logger.error(generateCodeFrame(sourceCode.length > 100_000 ? sourceCode : logger.highlight(nearest.file, sourceCode), 4, s))
       }
     })
   }
@@ -123,6 +125,7 @@ const skipErrorProperties = new Set([
   'type',
   'showDiff',
   'diff',
+  'codeFrame',
   'actual',
   'expected',
   'diffOptions',

packages/vitest/src/node/hoistMocks.ts

@@ -5,6 +5,7 @@ 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 { generateCodeFrame } from './error'
 
 export type Positioned<T> = T & {
@@ -256,7 +257,7 @@ export function hoistMocks(code: string, id: string, parse: PluginContext['parse
       name: 'SyntaxError',
       message: _error.message,
       stack: _error.stack,
-      frame: generateCodeFrame(code, 4, insideCall.start + 1),
+      frame: generateCodeFrame(highlight(code), 4, insideCall.start + 1),
     }
     throw error
   }

packages/vitest/src/node/logger.ts

@@ -1,5 +1,7 @@
 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'
@@ -21,19 +23,28 @@ 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
   errorStream = process.stderr
   logUpdate = createLogUpdate(process.stdout)
 
   private _clearScreenPending: string | undefined
+  private _highlights = new Map<string, string>()
 
   constructor(
     public ctx: Vitest,
     public console = globalThis.console,
   ) {
-
+    this._highlights.clear()
   }
 
   log(...args: any[]) {
@@ -91,6 +102,25 @@ export class Logger {
     })
   }
 
+  clearHighlightCache(filename?: string) {
+    if (filename)
+      this._highlights.delete(filename)
+    else
+      this._highlights.clear()
+  }
+
+  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 })
+    this._highlights.set(filename, code)
+    return code
+  }
+
   printNoTestFound(filters?: string[]) {
     const config = this.ctx.config
     const comma = c.dim(', ')

packages/vitest/src/runtime/mocker.ts

@@ -1,7 +1,7 @@
 import { existsSync, readdirSync } from 'node:fs'
 import vm from 'node:vm'
 import { basename, dirname, extname, isAbsolute, join, relative, resolve } from 'pathe'
-import { getColors, getType } from '@vitest/utils'
+import { getType, highlight } from '@vitest/utils'
 import { isNodeBuiltin } from 'vite-node/utils'
 import { distDir } from '../paths'
 import { getAllMockableProperties } from '../utils/base'
@@ -113,9 +113,11 @@ export class VitestMocker {
     return this.executor.state.filepath || 'global'
   }
 
-  private createError(message: string) {
+  private createError(message: string, codeFrame?: string) {
     const Error = this.primitives.Error
-    return new Error(message)
+    const error = new Error(message)
+    Object.assign(error, { codeFrame })
+    return error
   }
 
   public getMocks() {
@@ -208,18 +210,17 @@ export class VitestMocker {
         else if (!(prop in target)) {
           if (this.filterPublicKeys.includes(prop))
             return undefined
-          const c = getColors()
           throw this.createError(
             `[vitest] No "${String(prop)}" export is defined on the "${mockpath}" mock. `
             + 'Did you forget to return it from "vi.mock"?'
-            + '\nIf you need to partially mock a module, you can use "importOriginal" helper inside:\n\n'
-            + `${c.green(`vi.mock("${mockpath}", async (importOriginal) => {
+            + '\nIf you need to partially mock a module, you can use "importOriginal" helper inside:\n',
+            highlight(`vi.mock("${mockpath}", async (importOriginal) => {
   const actual = await importOriginal()
   return {
     ...actual,
     // your mocked methods
   }
-})`)}\n`,
+})`),
           )
         }