TypeScript：如何让按tag查找联合类型的函数返回正确类型？

我定义了如下联合类型：

type Foo = {tag: 'A', name: string} | {tag: 'B', value: number}

编写了findFirst函数用于查找数组中匹配指定tag的元素：

function findFirst(foos: Foo[], tag: 'A' | 'B') {
  for (const foo of foos) {
    if (foo.tag === tag) { 
      return foo 
    }
  }
  return undefined
}

但调用时，比如findFirst([{tag: 'A', name: 'X'}], 'A')?.name，TypeScript会报错Property 'name' does not exist on type 'Foo'——编译器无法根据传入的tag自动推断返回值的具体类型，需要显式定义函数返回类型。

解决方案 #

要让编译器关联tag参数和返回类型，需要使用泛型结合TypeScript的Extract工具类型来实现：

修改后的函数实现 #

type Foo = {tag: 'A', name: string} | {tag: 'B', value: number}

function findFirst<T extends Foo['tag']>(foos: Foo[], tag: T): Extract<Foo, {tag: T}> | undefined {
  for (const foo of foos) {
    if (foo.tag === tag) { 
      // 用类型断言告诉编译器当前foo的类型符合Extract筛选结果
      return foo as Extract<Foo, {tag: T}>
    }
  }
  return undefined
}

代码说明 #

1. 泛型T的作用：T extends Foo['tag']约束泛型T只能是Foo联合类型中存在的tag值（即'A'或'B'），这样编译器可以捕获调用时传入的具体tag类型。
2. Extract工具类型：Extract<Foo, {tag: T}>会从Foo联合类型中筛选出tag属性等于T的成员类型。比如当T为'A'时，返回类型就是{tag: 'A', name: string}；当T为'B'时，返回类型就是{tag: 'B', value: number}。
3. 类型断言：由于TypeScript的控制流分析暂时无法自动将foo.tag === tag的判断逻辑与泛型T关联，所以需要用as Extract<Foo, {tag: T}>做一次类型断言，确保编译器识别返回值的具体类型。

验证调用 #

现在调用函数时，编译器能正确推断返回值类型：

// 返回类型为 {tag: 'A', name: string} | undefined
const resultA = findFirst([{tag: 'A', name: 'X'}], 'A')
resultA?.name // 正常访问，无编译错误

// 返回类型为 {tag: 'B', value: number} | undefined
const resultB = findFirst([{tag: 'B', value: 123}], 'B')
resultB?.value // 正常访问，无编译错误

替代方案：自定义条件类型 #

如果不想使用Extract工具类型，也可以手动编写条件类型来关联tag和返回类型：

type Foo = {tag: 'A', name: string} | {tag: 'B', value: number}

// 自定义条件类型，根据tag返回对应的Foo成员类型
type FooByTag<T extends Foo['tag']> = 
  T extends 'A' ? {tag: 'A', name: string} : 
  T extends 'B' ? {tag: 'B', value: number} : 
  never

function findFirst<T extends Foo['tag']>(foos: Foo[], tag: T): FooByTag<T> | undefined {
  for (const foo of foos) {
    if (foo.tag === tag) { 
      return foo as FooByTag<T>
    }
  }
  return undefined
}

这种写法更直观，但当Foo联合类型的成员增多时，需要手动维护条件分支，而Extract可以自动适配Foo的定义，更适合复杂场景。

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