TypeScript: Infer type from discriminated union to create a map?

Question

enum ShapeType { Circle, Square }

interface Circle {
    kind: ShapeType.Circle,
    radius: number
}

interface Square {
    kind: ShapeType.Square,
    sideLength: number
}

type Shape = Circle | Square

type WhatShape<T extends ShapeType>   // <-- can we do without WhatShape?
    = T extends ShapeType.Circle ? Circle
    : T extends ShapeType.Square ? Square
    : never;

type ShapeMap = {
    [K in ShapeType]: (s: WhatShape<K>) => void
}


const handlers: ShapeMap = {
    [ShapeType.Circle]: (c: Circle) => {
        console.log(c.radius)
    },
    [ShapeType.Square]: s/*inferred!*/ => {
        console.log(s.sideLength)
    },
}

playground

This works, but it's a little annoying to have to maintain WhatShape to map the enum back to the shape kind.

Is there any way that TS can infer the argument type for our ShapeMap from the kind?

Answer 1

You can use the Extract<T, U> utility type to find the member(s) of a union T assignable to a type U. In your case, T would be the discriminated union type Shape, and U would be an object type whose kind property is the particular member(s) of ShapeType you are trying to map. So WhatShape<T> could be defined as:

type WhatShape<T extends ShapeType> = Extract<Shape, { kind: T }>

And it will behave the same as your existing version without needing to be maintained separately from Shape:

type ShapeMap = {
    [K in ShapeType]: (s: WhatShape<K>) => void
}
/* type ShapeMap = {
    0: (s: Circle) => void;
    1: (s: Square) => void;
} // same as before */

Playground link to code