【TypeScript 高级应用】函数如何根据参数类型动态推导出返回值类型

问题背景

假设现在有 Square 和 Circle 两个 interface，分别表示两种图形类型。

有一个 shapes 数组，里面装载的对象可以是 Square 也可以是 Circle。

getShapes(type) 方法用于返回 shapes 中某一特定类型的对象数组。

export interface Square {
  type: "square";
  width: number;
  height: number;
}

export interface Circle {
  type: "circle";
  radius: number;
}

const shapes: Array<Circle | Square> = [
  {
    type: "square",
    width: 10,
    height: 10,
  },
  {
    type: "circle",
    radius: 10,
  },
];

export function getShapes(type: Square["type"] | Circle["type"]) {
  return shapes.filter((item) => item.type === type);
}

上面的代码可以实现相关功能，但是 getShapes() 方法的返回类型推导结果是 Array<Square | Circle>：

const circles: Array<Square | Circle> = getShapes("circle");
const squares: Array<Square | Circle> = getShapes("square");

而我们期望的效果是让编译器能直接推导出结果为 Array<Square> 或 Array<Circle>：

const circles: Array<Circle> = getShapes("circle");
const squares: Array<Square> = getShapes("square");

函数重载

通过函数重载的方式，给 getShapes() 方法定义好传入参数类型对应的返回值类型，ts 就可以根据实际传入的参数类型来推导返回值的类型了。

export function getShapes(type: Square['type']): Array<Square>
export function getShapes(type: Circle['type']): Array<Circle>
export function getShapes(type: Square['type'] | Circle['type']) {
  return shapes.filter((item) => item.type === type)
}

const circles: Array<Circle> = getShapes("circle");
const squares: Array<Square> = getShapes("square");

泛型+类型推导+类型断言

将参数 type 的类型定义为泛型 T（T extends GetSharpsReq，只能为 circle 或 square 字符串），这样做的目的是为了让编译器在函数真正调用时才来确认这个 T 到底是什么。

将 shapes.filter() 的返回值用 as 设置为 GetSharpsRes<T>，而 GetSharpsRes<T> 内部进行了具体类型的推导。

当函数调用时 T 的类型被确定，这时编译器就可以推断出返回值的类型了。

type GetSharpsReq = Square['type'] | Circle['type']
type GetSharpsRes<T> = T extends Circle['type']
  ? Array<Circle>
  : T extends Square['type']
  ? Array<Square>
  : never

export function getShapes<T extends GetSharpsReq>(type: T) {
  return shapes.filter((item) => item.type === type) as GetSharpsRes<T>
}

const circles: Array<Circle> = getShapes('circle')
const squares: Array<Square> = getShapes('square')

参考

TS 中如何根据参数动态确定函数返回类型