Type function second argument dynamically based on first argument enum type

Question

Is it possible to infer the type of a second argument, which is a callback in my case, based on the enum value of the first argument.

I am having an issue where I am not able to infer the type for callback in the on function based on the eventName enum value.

interface User {
  name: string;
}
enum EventsEnum {
  JOIN = "join",
  LEAVE = "leave",
}
export interface EventsCallback {
  [EventsEnum.JOIN]: (users: User[]) => void;
  [EventsEnum.LEAVE]: (userId: string) => void;
}
export type Listener = {
  [eventName in EventsEnum]?: { callback: EventsCallback[eventName] };
};

class EventListener {
  listeners: Listener = {};

  on(eventName: EventsEnum, callback: EventsCallback[EventsEnum]) {
    if (!this.listeners[eventName]) {
      // I am having an issue here, this won't compile
      this.listeners[eventName] = { callback: callback };
    }
  }

  sampleDispatch(eventName: EventsEnum) {
    if (eventName === EventsEnum.JOIN) {
      this.listeners[eventName]?.callback([{ name: "user" }]);
    } else if (eventName === EventsEnum.LEAVE) {
      this.listeners[eventName]?.callback("user");
    }
  }
}

const event = new EventListener();
event.on(EventsEnum.JOIN, (users) => {
  //users here will infer correct type to User[]
});

event.on(EventsEnum.LEAVE, (user) => {
  //user here will infer correct type to string
});

Answer 1

The first problem you're having as that technically you can call on() with bad arguments. Your eventName is of the union type EventsEnum and your callback is of the completely independent union type EventsCallback[EventsEnum]. So you can call event.on(EventsEnum.JOIN, (user: string) => { }). You need to prevent that by making sure the type of eventName is correlated to the type of callback.

Now you could get that to happen with union types, but correlated union types aren't really supported directly, as described at microsoft/TypeScript#30581. The recommended approach is to refactor from unions to generic in the way described at microsoft/TypeScript#47109.

The call signature should look like this:

on<K extends EventsEnum>(eventName: K, callback: EventsMap[K]) {}

Now when you call event.on(), the generic type argument K is inferred from eventName, and that is used to give callback a specific type. So now the bad call is in error:

event.on(EventsEnum.JOIN, (user: string) => { }); // error!

And good calls give you the expected inference.

event.on(EventsEnum.JOIN, (users) => { });
//                         ^? (parameter) users: User[]

event.on(EventsEnum.LEAVE, (user) => { });
//                          ^? (parameter) user: string

---

Now callers are happy, but the implementation still gives you an error when you write this.listeners[eventName] = {callback: callback}. The compiler still can't follow the correlation between the generic eventName type K and the type of this.listeners. The way to fix this is to make Listener generic, and represent the operation this.listeners[eventName] as a distributive object type of the form {[P in K]: ⋯}[K]. Like this:

type Listener<K extends EventsEnum = EventsEnum> = {
  [P in K]?: { callback: EventsMap[P] };
}

on<K extends EventsEnum>(eventName: K, callback: EventsMap[K]) {
  const l: Listener<K> = this.listeners;
  if (!l[eventName]) {
    l[eventName] = { callback: callback };
  }
}

First, the Listener type is the same as before, but you can limit the keys you iterate over from EventsEnum to K. If you don't specify K, then you get the default type argument of EventsEnum, and so the type named Listener with no type argument is equivalent to the version from your code. That means you can still say that this.listeners is of type Listener.

Now, inside on(), you widen from Listener to Listener<K>. You can verify that Listener is assignable to Listener<K> for any K (since it's just possibly omitting some properties, and omitting properties is a valid widening). So I copy this.listeners to a variable l of type Listener<K>. Now l[eventName] is of the distributive object type { [P in K]?: { callback: EventsMap[P] }; }[K], which is a special form known to distribute over unions in K, and which the compiler will allow you to treat like {callback: EventsMap[K]} | undefined for assignment purposes (even though this is technically unsafe, see microsoft/TypeScript#48730). So l[eventName] = { callback: callback } is now allowed.

Playground link to code