如何为TypeScript通用Zip函数定义任意参数的泛型返回类型？

问题描述 #

我在TypeScript中实现了一个Zip组合函数，它可接收任意数量的数组作为参数，返回包含每个输入数组对应元素的新数组，但返回类型存在问题。该函数的返回类型需与Python或Lodash中的zip函数一致。当前实现代码如下：

function zipComp<T>(...arrs: T[][]) { ... }
zipComp([1, 2], [false, true]) // expected: [number, boolean][], actual: unknown[]
zipComp(["a", "b", "c"], [1, 2, 3], [{}, {}, {}]) // expected: [string, number, object][], actual: unknown[]

参考了Lodash的类型定义：

// from lodash (array.d.ts)
zip<T1, T2>(arrays1: List<T1>, arrays2: List<T2>): Array<[T1, T2]>;
zip<T1, T2, T3>(arrays1: List<T1>, arrays2: List<T2>, arrays3: List<T3>): Array<[T1, T2, T3]>;
zip<T1, T2, T3, T4>(arrays1: List<T1>, arrays2: List<T2>, arrays3: List<T3>, arrays4: List<T4>): Array<[T1, T2, T3, T4]>;
zip<T1, T2, T3, T4, T5>(arrays1: List<T1>, arrays2: List<T2>, arrays3: List<T3>, arrays4: List<T4>, arrays5: List<T5>): Array<[T1, T2, T3, T4, T5]>;
zip<T>(...arrays: Array<List<T>>): Array<Array<T>>; // => here is the problem

Lodash的实现仅在参数数量不超过5个时有效，超过则仍存在返回类型不准确的问题。请问是否存在无需逐个编写重载的通用参数类型定义方式？

解决方案 #

当然有，利用TypeScript的可变元组类型和映射类型可以实现完全通用的Zip函数类型，无需手动编写多个重载。

完整实现 #

type List<T> = ArrayLike<T>;

function zipComp<T extends List<any>[]>(...arrs: T): {
  [K in keyof T]: T[K] extends List<infer U> ? U : never
}[] {
  // 与Python zip行为一致：取最短数组的长度
  const minLength = Math.min(...arrs.map(arr => arr.length));
  return Array.from({ length: minLength }, (_, index) => 
    arrs.map(arr => arr[index])
  ) as ReturnType<typeof zipComp<T>>;
}

类型逻辑解析 #

• T extends List<any>[]：约束输入参数为任意数量的数组/类数组对象组成的元组，TypeScript会自动将传入的多个数组推断为具体的元组类型
• 映射类型{ [K in keyof T]: T[K] extends List<infer U> ? U : never }：遍历输入元组的每个元素，通过infer提取出每个数组的元素类型，最终组成对应位置类型的新元组
• 返回类型...[]：将上述元组类型包裹为数组，即得到每个位置对应元素类型的数组集合，和Python/Lodash的zip返回结构完全匹配

验证效果 #

// 推断类型：[number, boolean][]
const result1 = zipComp([1, 2], [false, true]);

// 推断类型：[string, number, object][]
const result2 = zipComp(["a", "b", "c"], [1, 2, 3], [{}, {}, {}]);

// 超过5个参数仍能准确推断
// 推断类型：[string, number, boolean, Date, RegExp, symbol][]
const result3 = zipComp(["x"], [1], [true], [new Date()], [/regex/], [Symbol()]);

注意点 #

• 实现中的类型断言as ReturnType<typeof zipComp<T>>是为了让TypeScript认可返回值的类型，对外暴露的类型是完全准确的
• 如果需要和Lodash一样取最长数组的长度（不足位置补undefined），只需把Math.min改成Math.max，类型会自动包含undefined（比如某数组长度不足时，对应位置类型会变为原类型 | undefined）

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