fix(types)!: do not try to infer types of overloaded functions (#2)

BREAKING CHANGE: overloaded function may now require explicit type annotations

Author: mcous

README.md

@@ -230,6 +230,27 @@ expect(spy('hello')).toEqual('goodbye');
 
 [asymmetric matchers]: https://vitest.dev/api/expect.html#expect-anything
 
+#### Types of overloaded functions
+
+Due to fundamental limitations in TypeScript, `when()` will always use the _last_ overload to infer function parameters and return types. You can use the `TFunc` type parameter of `when()` to manually select a different overload entry:
+
+```ts
+function overloaded(): null;
+function overloaded(input: number): string;
+function overloaded(input?: number): string | null {
+  // ...
+}
+
+// Last entry: all good!
+when(overloaded).calledWith(42).thenReturn('hello');
+
+// $ts-expect-error: first entry
+when(overloaded).calledWith().thenReturn(null);
+
+// Manually specified: all good!
+when<() => null>(overloaded).calledWith().thenReturn(null);
+```
+
 ### `.thenReturn(value: TReturn)`
 
 When the stubbing is satisfied, return `value`

src/behaviors.ts

@@ -1,22 +1,18 @@
 import { equals } from '@vitest/expect';
-import type {
-  AnyFunction,
-  AllParameters,
-  ReturnTypeFromArgs,
-} from './types.ts';
+import type { AnyFunction } from './types.ts';
 
 export const ONCE = Symbol('ONCE');
 
 export type StubValue<TValue> = TValue | typeof ONCE;
 
 export interface BehaviorStack<TFunc extends AnyFunction> {
   use: (
-    args: AllParameters<TFunc>
-  ) => BehaviorEntry<AllParameters<TFunc>> | undefined;
+    args: Parameters<TFunc>
+  ) => BehaviorEntry<Parameters<TFunc>> | undefined;
 
-  bindArgs: <TArgs extends AllParameters<TFunc>>(
+  bindArgs: <TArgs extends Parameters<TFunc>>(
     args: TArgs
-  ) => BoundBehaviorStack<ReturnTypeFromArgs<TFunc, TArgs>>;
+  ) => BoundBehaviorStack<ReturnType<TFunc>>;
 }
 
 export interface BoundBehaviorStack<TReturn> {
@@ -43,7 +39,7 @@ export interface BehaviorOptions<TValue> {
 export const createBehaviorStack = <
   TFunc extends AnyFunction
 >(): BehaviorStack<TFunc> => {
-  const behaviors: BehaviorEntry<AllParameters<TFunc>>[] = [];
+  const behaviors: BehaviorEntry<Parameters<TFunc>>[] = [];
 
   return {
     use: (args) => {

src/stubs.ts

@@ -1,12 +1,12 @@
 import type { Mock as Spy } from 'vitest';
 import { createBehaviorStack, type BehaviorStack } from './behaviors.ts';
 import { NotAMockFunctionError } from './errors.ts';
-import type { AnyFunction, AllParameters } from './types.ts';
+import type { AnyFunction } from './types.ts';
 
 const BEHAVIORS_KEY = Symbol('behaviors');
 
 interface WhenStubImplementation<TFunc extends AnyFunction> {
-  (...args: AllParameters<TFunc>): unknown;
+  (...args: Parameters<TFunc>): unknown;
   [BEHAVIORS_KEY]: BehaviorStack<TFunc>;
 }
 
@@ -25,7 +25,7 @@ export const configureStub = <TFunc extends AnyFunction>(
 
   const behaviors = createBehaviorStack<TFunc>();
 
-  const implementation = (...args: AllParameters<TFunc>): unknown => {
+  const implementation = (...args: Parameters<TFunc>): unknown => {
     const behavior = behaviors.use(args);
 
     if (behavior?.throwError) {
@@ -48,15 +48,15 @@ export const configureStub = <TFunc extends AnyFunction>(
 
 const validateSpy = <TFunc extends AnyFunction>(
   maybeSpy: unknown
-): Spy<AllParameters<TFunc>, unknown> => {
+): Spy<Parameters<TFunc>, unknown> => {
   if (
     typeof maybeSpy === 'function' &&
     'mockImplementation' in maybeSpy &&
     typeof maybeSpy.mockImplementation === 'function' &&
     'getMockImplementation' in maybeSpy &&
     typeof maybeSpy.getMockImplementation === 'function'
   ) {
-    return maybeSpy as Spy<AllParameters<TFunc>, unknown>;
+    return maybeSpy as Spy<Parameters<TFunc>, unknown>;
   }
 
   throw new NotAMockFunctionError(maybeSpy);

src/types.ts

@@ -1,60 +1,4 @@
-/**
- * Get function arguments and return value types.
- *
- * Support for overloaded functions, thanks to @Shakeskeyboarde
- * https://github.com/microsoft/TypeScript/issues/14107#issuecomment-1146738780
- */
-
-import type { SpyInstance } from 'vitest';
+/** Common type definitions. */
 
 /** Any function, for use in `extends` */
 export type AnyFunction = (...args: never[]) => unknown;
-
-/** Acceptable arguments for a function.*/
-export type AllParameters<TFunc extends AnyFunction> =
-  TFunc extends SpyInstance<infer TArgs, unknown>
-    ? TArgs
-    : Parameters<ToOverloads<TFunc>>;
-
-/** The return type of a function, given the actual arguments used.*/
-export type ReturnTypeFromArgs<
-  TFunc extends AnyFunction,
-  TArgs extends unknown[]
-> = TFunc extends SpyInstance<unknown[], infer TReturn>
-  ? TReturn
-  : ExtractReturn<ToOverloads<TFunc>, TArgs>;
-
-/** Given a functions and actual arguments used, extract the return type. */
-type ExtractReturn<
-  TFunc extends AnyFunction,
-  TArgs extends unknown[]
-> = TFunc extends (...args: infer TFuncArgs) => infer TFuncReturn
-  ? TArgs extends TFuncArgs
-    ? TFuncReturn
-    : never
-  : never;
-
-/** Transform an overloaded function into a union of functions. */
-type ToOverloads<TFunc extends AnyFunction> = Exclude<
-  OverloadUnion<(() => never) & TFunc>,
-  TFunc extends () => never ? never : () => never
->;
-
-/** Recursively extract functions from an overload into a union. */
-type OverloadUnion<TFunc, TPartialOverload = unknown> = TFunc extends (
-  ...args: infer TArgs
-) => infer TReturn
-  ? TPartialOverload extends TFunc
-    ? never
-    :
-        | OverloadUnion<
-            TPartialOverload & TFunc,
-            TPartialOverload &
-              ((...args: TArgs) => TReturn) &
-              OverloadProps<TFunc>
-          >
-        | ((...args: TArgs) => TReturn)
-  : never;
-
-/** Properties attached to a function. */
-type OverloadProps<TFunc> = Pick<TFunc, keyof TFunc>;

src/vitest-when.ts

@@ -1,18 +1,14 @@
 import { configureStub } from './stubs.ts';
 import type { StubValue } from './behaviors.ts';
-import type {
-  AnyFunction,
-  AllParameters,
-  ReturnTypeFromArgs,
-} from './types.ts';
+import type { AnyFunction } from './types.ts';
 
 export { ONCE, type StubValue } from './behaviors.ts';
 export * from './errors.ts';
 
 export interface StubWrapper<TFunc extends AnyFunction> {
-  calledWith<TArgs extends AllParameters<TFunc>>(
+  calledWith<TArgs extends Parameters<TFunc>>(
     ...args: TArgs
-  ): Stub<TArgs, ReturnTypeFromArgs<TFunc, TArgs>>;
+  ): Stub<TArgs, ReturnType<TFunc>>;
 }
 
 export interface Stub<TArgs extends unknown[], TReturn> {

test/typing.test-d.ts

@@ -31,15 +31,12 @@ describe('vitest-when type signatures', () => {
     assertType<subject.Stub<[1], string>>(stub);
   });
 
-  it('should reject invalid usage of a simple function', () => {
-    // @ts-expect-error: args missing
-    subject.when(simple).calledWith();
+  it('should handle a generic function', () => {
+    const stub = subject.when(generic).calledWith(1);
 
-    // @ts-expect-error: args wrong type
-    subject.when(simple).calledWith('hello');
+    stub.thenReturn('hello');
 
-    // @ts-expect-error: return wrong type
-    subject.when(simple).calledWith(1).thenReturn(42);
+    assertType<subject.Stub<[number], string>>(stub);
   });
 
   it('should handle an overloaded function using its last overload', () => {
@@ -50,30 +47,14 @@ describe('vitest-when type signatures', () => {
     assertType<subject.Stub<[1], string>>(stub);
   });
 
-  it('should handle an overloaded function using its first overload', () => {
-    const stub = subject.when(overloaded).calledWith();
+  it('should handle an overloaded function using an explicit type', () => {
+    const stub = subject.when<() => null>(overloaded).calledWith();
 
     stub.thenReturn(null);
 
     assertType<subject.Stub<[], null>>(stub);
   });
 
-  it('should handle an very overloaded function using its first overload', () => {
-    const stub = subject.when(veryOverloaded).calledWith();
-
-    stub.thenReturn(null);
-
-    assertType<subject.Stub<[], null>>(stub);
-  });
-
-  it('should handle an overloaded function using its last overload', () => {
-    const stub = subject.when(veryOverloaded).calledWith(1, 2, 3, 4);
-
-    stub.thenReturn(42);
-
-    assertType<subject.Stub<[1, 2, 3, 4], number>>(stub);
-  });
-
   it('should reject invalid usage of a simple function', () => {
     // @ts-expect-error: args missing
     subject.when(simple).calledWith();
@@ -84,6 +65,17 @@ describe('vitest-when type signatures', () => {
     // @ts-expect-error: return wrong type
     subject.when(simple).calledWith(1).thenReturn(42);
   });
+
+  it('should reject invalid usage of a generic function', () => {
+    // @ts-expect-error: args missing
+    subject.when(generic).calledWith();
+
+    // @ts-expect-error: args wrong type
+    subject.when(generic<string>).calledWith(42);
+
+    // @ts-expect-error: return wrong type
+    subject.when(generic).calledWith(1).thenReturn(42);
+  });
 });
 
 function untyped(...args: any[]): any {
@@ -94,22 +86,12 @@ function simple(input: number): string {
   throw new Error(`simple(${input})`);
 }
 
+function generic<T>(input: T): string {
+  throw new Error(`generic(${input})`);
+}
+
 function overloaded(): null;
 function overloaded(input: number): string;
 function overloaded(input?: number): string | null {
   throw new Error(`overloaded(${input})`);
 }
-
-function veryOverloaded(): null;
-function veryOverloaded(i1: number): string;
-function veryOverloaded(i1: number, i2: number): boolean;
-function veryOverloaded(i1: number, i2: number, i3: number): null;
-function veryOverloaded(i1: number, i2: number, i3: number, i4: number): number;
-function veryOverloaded(
-  i1?: number,
-  i2?: number,
-  i3?: number,
-  i4?: number
-): string | boolean | number | null {
-  throw new Error(`veryOverloaded(${i1}, ${i2}, ${i3}, ${i4})`);
-}