如何让函数从联合类型变量返回对应联合类型？

我需要从一个对象联合类型的源数据创建新对象，但TypeScript无法将新对象的类型推断为对应的区分联合类型。当前函数返回类型为{type: "ws" | "file"; state: "connecting" | number;}，而非预期的{type: "file"; state: number;} | {type: "ws"; state: "connecting";}，导致类型收窄时出现错误。我不想使用switch case或手动指定返回类型，希望代码具备模块化，新增源类型时无需修改相关代码。

问题代码示例 #

type Source = {
    type: 'ws'
    reader: {state: 'connecting'}
} | {
    type: 'file'
    reader: {state: number}
}

let source: Source = {type: 'file', reader:  {state: 12}}

/*
 * 错误：getSourceWithState 返回类型是 
 * {
 *   type: "ws" | "file";
 *   state: "connecting" | number;
 * } 
 * 而非预期的
 *  {
 *   type: "file";
 *   state: number;
 * } |
 *  {
 *   type: "ws";
 *   state: "connecting";
 * } 
 */
const getSourceWithState = () => {
    return {
        type: source.type,
        state: source.reader.state
    }
}

const sourceWithState = getSourceWithState()
if (sourceWithState.type === 'file') {
    // 错误：无法将类型“"connecting" | number”分配给类型“number”
    sourceWithState.state += 1
}

解决方案：泛型+映射类型自动推导关联类型 #

通过泛型约束和映射类型，可以让TypeScript正确保留联合类型中各成员的类型关联，无需手动维护返回类型或使用switch case，新增源类型时也无需修改核心代码：

type Source = {
    type: 'ws'
    reader: {state: 'connecting'}
} | {
    type: 'file'
    reader: {state: number}
}

// 定义映射类型，将Source的每个分支转换为目标结构
type SourceWithState<S extends Source> = {
    type: S['type']
    state: S['reader']['state']
}

// 泛型函数，接收具体的Source分支并返回对应的关联类型
const getSourceWithState = <S extends Source>(src: S): SourceWithState<S> => {
    return {
        type: src.type,
        state: src.reader.state
    } as SourceWithState<S>
}

let source: Source = {type: 'file', reader:  {state: 12}}
const sourceWithState = getSourceWithState(source)

if (sourceWithState.type === 'file') {
    // 类型收窄正常，可以安全操作number类型的state
    sourceWithState.state += 1
}

方案说明 #

• 泛型<S extends Source>让函数参数被推断为联合类型中的具体分支，而非整个联合类型，保留了类型内部的关联关系
• SourceWithState映射类型针对每个具体分支生成对应的结构，确保type和state的类型绑定不丢失
• 新增Source的类型分支时，只需扩展Source联合类型，getSourceWithState函数会自动适配新的类型，无需修改现有代码

内容的提问来源于Stack Exchange，提问作者Ybri