如何根据keyof RootState的字符串参数自动推导函数返回类型？

现有如下TypeScript类型定义与状态对象：

type Car = {
    name: string,
    brand: string,
}
type User = {
    name: string, money: number
}
type Driver = {
    name: string, age: string
}
type RootState = {
    car: Car[],    
    user: User[],
    driver: Driver[],
}

const state: RootState = {
    car: [{name: "Car1", brand: "Ford"}, {name: "Car2", brand: "Chevvy"}],
    driver: [{name: "joe", age: "3"}, {name: "carl", age: "4"}],
    user: [{name: "Hank", money: 1300}]
}

希望编写一个函数，接收RootState的键名作为第一个参数，自动推导返回对应数组的元素类型（比如传入"car"时返回Car类型），但当前实现的函数返回类型是Car | User | Driver联合类型，不符合预期：

function getResourceAtIndex(type: keyof RootState, index: number){
    return state[type][index]
}

const car = getResourceAtIndex("car", 0) // 返回类型为 Car | User | Driver，而非预期的 Car

修改方案 #

要让TypeScript根据传入的键名精准推导返回类型，需要给函数添加泛型约束，将参数类型与RootState的键绑定，明确返回类型为对应数组的元素类型：

修改后的函数实现：

function getResourceAtIndex<T extends keyof RootState>(type: T, index: number): RootState[T][number] {
    return state[type][index]
}

也可以省略显式返回类型，TypeScript会自动完成推导：

function getResourceAtIndex<T extends keyof RootState>(type: T, index: number) {
    return state[type][index]
}

验证效果：

const car = getResourceAtIndex("car", 0) // 类型自动推导为 Car
const user = getResourceAtIndex("user", 0) // 类型自动推导为 User
const driver = getResourceAtIndex("driver", 0) // 类型自动推导为 Driver

原理说明 #

• 泛型T extends keyof RootState限定了type参数只能是RootState的合法键名，建立了参数与类型的关联
• RootState[T][number]表示取RootState中键T对应数组的元素类型，让编译器能根据传入的具体键名，精准推导返回值的类型

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