如何实现多键分组求和的TypeScript类型安全通用函数？

需求说明 #

给定对象数组，需要按任意数量的对象键进行分组，并对另一组任意数量的键的值求和。例如：

输入数组：

const arr = [
   { shape: "square", color: "red", available: 1, ordered: 1 },
   { shape: "square", color: "red", available: 2, ordered: 1 },
   { shape: "circle", color: "blue", available: 0, ordered: 3 },
   { shape: "square", color: "blue", available: 4, ordered: 4 },
];

按shape和color分组，对available和ordered求和后，预期输出：

[
  { shape: "square", color: "red", available: 3, ordered: 2 },
  { shape: "circle", color: "blue", available: 0, ordered: 3 },
  { shape: "square", color: "blue", available: 4, ordered: 4 },
];

现有方案存在无通用可用函数、缺乏TypeScript类型安全、无法妥善处理对象额外键的问题，因此需要构建支持多分组键与多求和键的通用groupBySum函数。

完整实现 #

类型定义与函数实现 #

type NumericKeys<T> = { [K in keyof T]: T[K] extends number ? K : never }[keyof T];

function groupBySum<
  T extends Record<string, any>,
  GroupKeys extends keyof T,
  SumKeys extends NumericKeys<T>
>(
  arr: T[],
  groupKeys: GroupKeys[],
  sumKeys: SumKeys[]
): Array<Pick<T, GroupKeys> & Record<SumKeys, number>> {
  const groupMap = new Map<string, Pick<T, GroupKeys> & Record<SumKeys, number>>();

  for (const item of arr) {
    // 生成分组唯一标识
    const groupKey = groupKeys.map(key => item[key]).join('|');
    
    // 获取或创建分组对象
    let groupItem = groupMap.get(groupKey);
    if (!groupItem) {
      // 初始化分组键的值
      const groupBase = groupKeys.reduce((acc, key) => {
        acc[key] = item[key];
        return acc;
      }, {} as Pick<T, GroupKeys>);
      // 初始化求和键为0
      const sumBase = sumKeys.reduce((acc, key) => {
        acc[key] = 0;
        return acc;
      }, {} as Record<SumKeys, number>);
      
      groupItem = { ...groupBase, ...sumBase };
      groupMap.set(groupKey, groupItem);
    }

    // 累加求和键的值
    for (const key of sumKeys) {
      groupItem[key] += item[key] ?? 0;
    }
  }

  return Array.from(groupMap.values());
}

代码说明 #

1. 类型安全保障：

   • NumericKeys<T> 提取对象中值为数字类型的键，确保求和键只能是数字类型的属性
   • 泛型约束确保groupKeys是输入对象的键，sumKeys是输入对象的数字类型键，调用时传错键会直接触发TypeScript报错

2. 分组逻辑：

   • 使用Map存储分组，以分组键值的拼接字符串作为唯一键（避免对象作为Map键的引用问题）
   • 首次遇到分组时，复制分组键的原始值，并初始化求和键为0
   • 后续同分组的元素直接累加求和键的值

3. 额外键处理：

   • 函数返回的是Pick<T, GroupKeys> & Record<SumKeys, number>类型，仅保留分组键和求和键，不会包含原对象的额外键（若需要保留额外键，可调整返回类型为T，但需确保额外键在同分组中值一致，否则会取第一个元素的值）

示例调用 #

const result = groupBySum(arr, ['shape', 'color'], ['available', 'ordered']);
console.log(result);
// 输出与预期一致

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