如何基于函数参数推断TypeScript函数的返回类型？

问题场景 #

现有如下TypeScript代码：

const handleNodes = (node: Node | Node[]) => {
  if (Array.isArray(node)) {
    return [{}];
  }
  return {};
};

期望实现：调用handleNodes([{}])时自动推断返回数组类型{}[]，调用handleNodes({})时自动推断返回对象类型{}，尝试条件类型后未成功，需找到正确实现方式。

两种可行方案 #

方案一：泛型结合条件类型 #

通过泛型绑定输入参数的类型，配合条件类型明确返回值与输入的对应关系：

// 定义条件类型，根据输入是否为数组决定返回类型
type HandleNodesReturn<T> = T extends Node[] ? {}[] : {};

const handleNodes = <T extends Node | Node[]>(node: T): HandleNodesReturn<T> => {
  if (Array.isArray(node)) {
    // 用类型断言告知编译器返回值符合条件类型
    return [{}] as HandleNodesReturn<T>;
  }
  return {} as HandleNodesReturn<T>;
};

方案二：函数重载 #

通过函数重载直接定义不同输入对应的返回类型，写法更直观：

// 重载签名：输入为数组时返回数组
function handleNodes(node: Node[]): {}[];
// 重载签名：输入为单个节点时返回对象
function handleNodes(node: Node): {};
// 实现签名：包含所有输入情况的逻辑
function handleNodes(node: Node | Node[]) {
  if (Array.isArray(node)) {
    return [{}];
  }
  return {};
}

效果验证 #

• 调用handleNodes([{}])时，TypeScript会自动推断返回类型为{}[]
• 调用handleNodes({})时，自动推断返回类型为{}

内容的提问来源于stack exchange，提问作者jsardev