如何为TypeScript消息处理器对象实现与switch一致的类型收窄？

原本使用switch/case处理WebSocket消息时，TypeScript能自动对message参数进行类型收窄，但切换为处理器对象映射的形式后，无法让处理器的参数自动收窄到对应消息类型。

原switch/case实现示例：

const handleMessage = (message: Message) => {
  switch (message.type) {
    case 'heartbeat': {
      if (message.action === 'ping') {
        // message自动收窄为heartbeat-ping类型
        console.log("Heartbeat", message.timestamp);
        break;
      }
      // ...其他heartbeat action处理
    }
    case 'chat': {
      // ...chat类型处理
    }
  }
};

想要切换的处理器对象形式：

const handlers = {
  heartbeat: {
    ping: (msg) => console.log(msg)
  },
  // ...其他类型和action处理器
};

const handleMessage2 = (message: Message) => {
  handlers[message.type]?.[message.action]?.(message);
};

1. 明确消息联合类型 #

首先确保你的Message是区分联合类型（通过type和action字段区分）：

type Message = 
  | { type: 'heartbeat', action: 'ping', timestamp: number }
  | { type: 'heartbeat', action: 'pong', timestamp: number }
  | { type: 'chat', action: 'send', content: string, userId: string }
  | { type: 'chat', action: 'receive', content: string, senderId: string };

2. 构建处理器映射类型 #

基于Message类型，自动生成严格的处理器类型定义：

// 为每个消息类型+action组合生成对应的处理器类型
type MessageHandlers = {
  [M in Message as M['type']]: {
    [A in M['action']]: (msg: M) => void;
  };
};

3. 定义类型安全的处理器对象 #

此时TypeScript会自动校验每个处理器的参数类型，并实现自动收窄：

const handlers: MessageHandlers = {
  heartbeat: {
    ping: (msg) => {
      // msg自动收窄为{ type: 'heartbeat', action: 'ping', timestamp: number }
      console.log("Heartbeat ping", msg.timestamp);
    },
    pong: (msg) => {
      console.log("Heartbeat pong", msg.timestamp);
    }
  },
  chat: {
    send: (msg) => {
      // msg自动收窄为{ type: 'chat', action: 'send', content: string, userId: string }
      console.log("Chat send from", msg.userId, ":", msg.content);
    },
    receive: (msg) => {
      console.log("Chat receive from", msg.senderId, ":", msg.content);
    }
  }
};

4. 实现类型安全的消息处理函数 #

通过类型断言确保调用处理器时的类型匹配：

const handleMessage = (message: Message) => {
  // 确保action是当前type对应的合法action
  const actionHandler = handlers[message.type]?.[message.action as keyof typeof handlers[typeof message.type]];
  if (actionHandler) {
    (actionHandler as (msg: typeof message) => void)(message);
  }
};

或者更简洁的可选链写法：

const handleMessage = (message: Message) => {
  (handlers[message.type]?.[message.action as keyof typeof handlers[typeof message.type]] as (msg: typeof message) => void)?.(message);
};

原理说明 #

MessageHandlers利用TypeScript的映射类型和条件类型，遍历Message联合类型的每个子类型，为每个type下的action生成精确的处理器签名。这样在定义handlers时，TypeScript会自动校验处理器参数类型，调用时也能确保传入的message与处理器的参数类型完全匹配。

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