如何通过函数参数收窄TypeScript可辨识联合类型

我有一个通过ID存储多种不同宠物类型的状态对象，定义如下：

interface Dog { type: "dog"; woofs: string; }
interface Cat { type: "cat"; meows: string; }
type Pet = Dog | Cat;
type AnimalState = Record<string, Pet>

我希望实现一个函数，通过ID和指定类型返回特定的宠物类型：

function getSpecificAnimal(state: AnimalState, key: string, type: Pet["type"]) {
  const pet = state[key];

  if (pet.type === type) {
    return pet;
  }

  throw new Error("pet was wrong type");
}

const aDog = getSpecificAnimal(state, "a", "dog"); // 应返回Dog类型或抛出异常

但我无法让TypeScript正确收窄getSpecificAnimal函数的返回类型。我尝试使用泛型和类型映射关联参数，但TypeScript仍不认可：

type PetTypeMap = {
  dog: Dog,
  cat: Cat,
}
function getAnimalOfType<T extends Pet, K extends T["type"]>(animals: AnimalState, id: string, type: K): PetTypeMap[K] {
  const pet = animals[id];
  //    ^?

  if (pet.type === type) {
    return pet;
  }

  throw new Error("wrong pet type");
}

我还尝试了多种其他方法，但都无法解决问题。请问在TypeScript中是否无法实现这类类型收窄？

可以实现这类类型收窄，问题出在泛型约束逻辑和类型收窄的辅助手段上，以下是两种可行的实现方式：

方式一：泛型+类型断言 #

首先明确PetTypeMap的类型映射关系，将泛型约束为keyof PetTypeMap，让TypeScript直接关联传入的类型参数和返回类型，再通过类型断言解决收窄问题：

interface Dog { type: "dog"; woofs: string; }
interface Cat { type: "cat"; meows: string; }
type Pet = Dog | Cat;
type AnimalState = Record<string, Pet>

type PetTypeMap = {
  dog: Dog,
  cat: Cat,
};

function getSpecificAnimal<K extends keyof PetTypeMap>(
  state: AnimalState,
  key: string,
  type: K
): PetTypeMap[K] {
  const pet = state[key];
  
  if (pet.type === type) {
    return pet as PetTypeMap[K];
  }

  throw new Error(`pet was wrong type, expected ${type}`);
}

// 使用示例
const state: AnimalState = {
  a: { type: "dog", woofs: "loud" },
  b: { type: "cat", meows: "soft" }
};

const aDog = getSpecificAnimal(state, "a", "dog"); // 类型自动推断为Dog
const aCat = getSpecificAnimal(state, "b", "cat"); // 类型自动推断为Cat

方式二：自定义类型守卫（更严谨） #

如果想让TypeScript在编译阶段自动完成类型收窄，无需手动断言，可以自定义类型守卫函数：

interface Dog { type: "dog"; woofs: string; }
interface Cat { type: "cat"; meows: string; }
type Pet = Dog | Cat;
type AnimalState = Record<string, Pet>

type PetTypeMap = {
  dog: Dog,
  cat: Cat,
};

// 自定义类型守卫，告诉TypeScript如何判断pet是否为指定类型
function isPetOfType<K extends keyof PetTypeMap>(
  pet: Pet,
  type: K
): pet is PetTypeMap[K] {
  return pet.type === type;
}

function getSpecificAnimal<K extends keyof PetTypeMap>(
  state: AnimalState,
  key: string,
  type: K
): PetTypeMap[K] {
  const pet = state[key];
  
  if (isPetOfType(pet, type)) {
    return pet; // 这里TypeScript会自动将pet收窄为PetTypeMap[K]类型
  }

  throw new Error(`pet was wrong type, expected ${type}`);
}

失败原因分析 #

• 之前的泛型定义T extends Pet, K extends T["type"]没有建立K与返回类型的直接绑定，TypeScript无法推断出pet的具体类型。
• 直接的pet.type === type在泛型场景下，TypeScript无法自动完成联合类型的收窄，需要类型断言或自定义类型守卫来辅助类型推断。

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