TypeScript中如何让函数返回值匹配参数的键联合类型

我有如下TypeScript函数：

function doSomething(data: { [key: string]: [number, number] }, height: number): { matchingData: string | undefined } {
  for (const [key, [min, max]] of Object.entries(data)) {
    if (height >= min && height < max) {
      return { matchingData: key };
    }
  }

  return { matchingData: undefined };
}

我希望该函数返回值的matchingData字段类型为入参data的键的字符串联合类型，这样在使用函数时能获得语法补全。比如下面的示例中，直接判断时没有补全，但手动断言类型后就有了：

const data = doSomething({ small: [0, 50], medium: [50, 100], large: [100, 150] }, 42);
if (data.matchingData === "small") {
  console.log("no syntax completion");
}

if (data.matchingData as "small" | "medium" | "large" === "small") {
  console.log("has syntax completion");
}

我尝试过keyof等方法但未成功，想请教该如何解决，是否有更合适的入参定义方式？

核心是用泛型捕获入参data的键类型，让TypeScript自动推断返回值的具体联合类型，具体修改如下：

修改后的函数代码 #

function doSomething<K extends string>(
  data: Record<K, [number, number]>,
  height: number
): { matchingData: K | undefined } {
  for (const [key, [min, max]] of Object.entries(data)) {
    if (height >= min && height < max) {
      // 类型断言：Object.entries返回的key默认是string，这里明确它就是K类型
      return { matchingData: key as K };
    }
  }
  return { matchingData: undefined };
}

关键说明 #

1. 泛型K：K extends string限定泛型为字符串子集，用来捕获传入data的所有键的具体类型。
2. 入参类型：Record<K, [number, number]>是TypeScript提供的快捷语法，等价于{ [key in K]: [number, number] }，表示data的键类型为K，值类型为[number, number]。
3. 类型断言的必要性：Object.entries()返回的键默认被推断为string，但我们明确知道它属于K类型，因此这个断言逻辑安全。

使用效果 #

现在调用函数时，返回值的matchingData会自动推断为传入data的键的联合类型，编辑器直接提供语法补全：

const result = doSomething({ small: [0, 50], medium: [50, 100], large: [100, 150] }, 42);
if (result.matchingData === "small") {
  console.log("现在有语法补全了");
}
// 此时result.matchingData的类型为 "small" | "medium" | "large" | undefined

内容的提问来源于stack exchange，提问作者julian.a