如何让TypeScript根据标签值推断标记联合的返回子类型？

问题概述 #

已知Foo类型对象包含bar: Bar成员，其中Bar是标记联合类型：

type ABar = {name: 'A', aData: string};
type BBar = {name: 'B', bData: string};
type Bar = ABar | BBar;
type BarName = Bar['name']; // 等价于 'A'|'B'

现有提取函数getBar需要根据传入的name返回对应类型的Bar实例，但当前实现要求调用者手动断言类型，无法让TypeScript自动推断返回值类型：

function getBar(foo: Foo, name: BarName): Bar | null {
    return foo.bar.name === name ? foo.bar : null;
}

// 必须显式断言类型
let aBar = getBar(foo, 'A') as ABar | null;

需要实现：当传入name='A'时，函数返回ABar | null；传入name='B'时返回BBar | null，且无需手动断言。

解决方案 #

方案1：利用类型映射表 #

1. 先定义BarName到对应子类型的映射：

type BarTypeMap = {
    'A': ABar;
    'B': BBar;
};

2. 改造为泛型函数，通过泛型参数约束name类型，并映射返回值：

type Foo = { bar: Bar };

function getBar<N extends BarName>(foo: Foo, name: N): BarTypeMap[N] | null {
    return foo.bar.name === name ? foo.bar as BarTypeMap[N] : null;
}

方案2：直接使用Extract条件类型（更简洁） #

无需额外定义映射表，直接用TypeScript内置的Extract工具类型从联合类型中筛选目标子类型：

type Foo = { bar: Bar };

function getBar<N extends BarName>(foo: Foo, name: N): Extract<Bar, { name: N }> | null {
    return foo.bar.name === name ? foo.bar as Extract<Bar, { name: N }> : null;
}

使用效果 #

两种方案都能让TypeScript自动推断返回值类型：

const foo: Foo = { bar: { name: 'A', aData: 'demo' } };

// 自动推断为 ABar | null
const aBar = getBar(foo, 'A');
// 自动推断为 BBar | null
const bBar = getBar(foo, 'B');

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