From a749a6c01fb6a577de50a78ffcc91014c2a6fe55 Mon Sep 17 00:00:00 2001
From: Han <fenghan770@gmail.com>
Date: Mon, 3 Jul 2023 15:59:47 +0800
Subject: [PATCH] feat: support use function/class as `bench` name (#3711)

---
 docs/api/index.md                                 |  8 ++++----
 packages/vitest/src/runtime/benchmark.ts          |  8 ++++++--
 packages/vitest/src/types/benchmark.ts            |  2 +-
 test/reporters/fixtures/function-as-name.bench.ts |  8 ++++++++
 test/reporters/tests/function-as-name.test.ts     | 10 ++++++++++
 5 files changed, 29 insertions(+), 7 deletions(-)
 create mode 100644 test/reporters/fixtures/function-as-name.bench.ts

diff --git a/docs/api/index.md b/docs/api/index.md
index a9d43126a023..8c964a9b6dbc 100644
--- a/docs/api/index.md
+++ b/docs/api/index.md
@@ -344,7 +344,7 @@ You cannot use this syntax, when using Vitest as [type checker](/guide/testing-t
 
 ## bench
 
-- **Type:** `(name: string, fn: BenchFunction, options?: BenchOptions) => void`
+- **Type:** `(name: string | Function, fn: BenchFunction, options?: BenchOptions) => void`
 
 `bench` defines a benchmark. In Vitest terms benchmark is a function that defines a series of operations. Vitest runs this function multiple times to display different performance results.
 
@@ -411,7 +411,7 @@ Vitest uses [`tinybench`](https://github.com/tinylibs/tinybench) library under t
 
 ### bench.skip
 
-- **Type:** `(name: string, fn: BenchFunction, options?: BenchOptions) => void`
+- **Type:** `(name: string | Function, fn: BenchFunction, options?: BenchOptions) => void`
 
 You can use `bench.skip` syntax to skip running certain benchmarks.
 
@@ -428,7 +428,7 @@ You can use `bench.skip` syntax to skip running certain benchmarks.
 
 ### bench.only
 
-- **Type:** `(name: string, fn: BenchFunction, options?: BenchOptions) => void`
+- **Type:** `(name: string | Function, fn: BenchFunction, options?: BenchOptions) => void`
 
 Use `bench.only` to only run certain benchmarks in a given suite. This is useful when debugging.
 
@@ -445,7 +445,7 @@ Use `bench.only` to only run certain benchmarks in a given suite. This is useful
 
 ### bench.todo
 
-- **Type:** `(name: string) => void`
+- **Type:** `(name: string | Function) => void`
 
 Use `bench.todo` to stub benchmarks to be implemented later.
 
diff --git a/packages/vitest/src/runtime/benchmark.ts b/packages/vitest/src/runtime/benchmark.ts
index 0e25180e474d..6c1afac7b7d2 100644
--- a/packages/vitest/src/runtime/benchmark.ts
+++ b/packages/vitest/src/runtime/benchmark.ts
@@ -21,7 +21,7 @@ export const bench = createBenchmark(
     if (!isRunningInBenchmark())
       throw new Error('`bench()` is only available in benchmark mode.')
 
-    const task = getCurrentSuite().custom.call(this, name)
+    const task = getCurrentSuite().custom.call(this, formatName(name))
     task.meta = {
       benchmark: true,
     }
@@ -33,7 +33,7 @@ export const bench = createBenchmark(
 function createBenchmark(fn: (
   (
     this: Record<'skip' | 'only' | 'todo', boolean | undefined>,
-    name: string,
+    name: string | Function,
     fn?: BenchFunction,
     options?: BenchOptions
   ) => void
@@ -48,3 +48,7 @@ function createBenchmark(fn: (
 
   return benchmark as BenchmarkAPI
 }
+
+function formatName(name: string | Function) {
+  return typeof name === 'string' ? name : name instanceof Function ? (name.name || '<anonymous>') : String(name)
+}
diff --git a/packages/vitest/src/types/benchmark.ts b/packages/vitest/src/types/benchmark.ts
index e0b3368dd95f..74cc58ebf473 100644
--- a/packages/vitest/src/types/benchmark.ts
+++ b/packages/vitest/src/types/benchmark.ts
@@ -54,7 +54,7 @@ export interface BenchmarkResult extends TinybenchResult {
 export type BenchFunction = (this: BenchFactory) => Promise<void> | void
 export type BenchmarkAPI = ChainableFunction<
 'skip' | 'only' | 'todo',
-[name: string, fn?: BenchFunction, options?: BenchOptions],
+[name: string | Function, fn?: BenchFunction, options?: BenchOptions],
 void
 > & {
   skipIf(condition: any): BenchmarkAPI
diff --git a/test/reporters/fixtures/function-as-name.bench.ts b/test/reporters/fixtures/function-as-name.bench.ts
new file mode 100644
index 000000000000..c7291be53eb8
--- /dev/null
+++ b/test/reporters/fixtures/function-as-name.bench.ts
@@ -0,0 +1,8 @@
+import { bench } from 'vitest'
+
+function foo() {}
+class Bar {}
+
+bench(foo, () => {})
+bench(Bar, () => {})
+bench(() => {}, () => {})
diff --git a/test/reporters/tests/function-as-name.test.ts b/test/reporters/tests/function-as-name.test.ts
index 2dda78176898..15c3273f5eba 100644
--- a/test/reporters/tests/function-as-name.test.ts
+++ b/test/reporters/tests/function-as-name.test.ts
@@ -14,3 +14,13 @@ test('should print function name', async () => {
   expect(stdout).toContain('function-as-name.test.ts > foo > foo')
   expect(stdout).toContain('function-as-name.test.ts > Bar > Bar')
 })
+
+test('should print function name in benchmark', async () => {
+  const filename = resolve('./fixtures/function-as-name.bench.ts')
+  const { stdout } = await runVitest({ root: './fixtures' }, [filename], 'benchmark')
+
+  expect(stdout).toBeTruthy()
+  expect(stdout).toContain('Bar')
+  expect(stdout).toContain('foo')
+  expect(stdout).toContain('<anonymous>')
+})