/
benchmark.ts
54 lines (45 loc) · 1.66 KB
/
benchmark.ts
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
import type { TaskCustom } from '@vitest/runner'
import { getCurrentSuite } from '@vitest/runner'
import { createChainable } from '@vitest/runner/utils'
import { noop } from '@vitest/utils'
import type { BenchFunction, BenchOptions, BenchmarkAPI } from '../types'
import { isRunningInBenchmark } from '../utils'
const benchFns = new WeakMap<TaskCustom, BenchFunction>()
const benchOptsMap = new WeakMap()
export function getBenchOptions(key: TaskCustom): BenchOptions {
return benchOptsMap.get(key)
}
export function getBenchFn(key: TaskCustom): BenchFunction {
return benchFns.get(key)!
}
export const bench = createBenchmark(
function (name, fn: BenchFunction = noop, options: BenchOptions = {}) {
if (!isRunningInBenchmark())
throw new Error('`bench()` is only available in benchmark mode.')
const task = getCurrentSuite().custom.call(this, formatName(name))
task.meta = {
benchmark: true,
}
benchFns.set(task, fn)
benchOptsMap.set(task, options)
},
)
function createBenchmark(fn: (
(
this: Record<'skip' | 'only' | 'todo', boolean | undefined>,
name: string | Function,
fn?: BenchFunction,
options?: BenchOptions
) => void
)) {
const benchmark = createChainable(
['skip', 'only', 'todo'],
fn,
) as BenchmarkAPI
benchmark.skipIf = (condition: any) => (condition ? benchmark.skip : benchmark) as BenchmarkAPI
benchmark.runIf = (condition: any) => (condition ? benchmark : benchmark.skip) as BenchmarkAPI
return benchmark as BenchmarkAPI
}
function formatName(name: string | Function) {
return typeof name === 'string' ? name : name instanceof Function ? (name.name || '<anonymous>') : String(name)
}