How can I map discriminated unions to concrete types (like classes) via a factory function in Typescript?

Question

I'm trying something like the following:

type Node =
    | { type: 'Group'; children: Node[] }
    | { type: 'Obj'; payload: number };

class Group {
    readonly type = 'Group';
    constructor(public n: Node[]) { }
}

class Obj {
    readonly type = 'Obj';
    constructor(public p: number) { }
}

type NodeMappings =
    { Group: Group; Obj: Obj };

function createNode<T extends Node>(node: T): NodeMappings[T['type']] {
    switch (node.type) {
        case 'Group':
            return new Group(node.children); // error!
        case 'Obj':
            return new Obj(node.payload); // error!
    }
}

Is there a way to achieve something like the given code that works with all kinds of types for the mapping in the createNode function (here done via a switch ... case)? Is this a bug in Typescript?

Answer 1

The compiler cannot just see the abstract relationship between the Node and NodeMappings in order to verify that createNode() is properly implemented.

If you want to express that relationship in a way the compiler understands, you need to do so explicitly by following the steps laid out in microsoft/TypeScript#47109. Your need a "base" type which looks like a key-value mapping from each type property to the rest of the relevant Node member:

type NodeMap = { [N in Node as N['type']]:
    { [P in keyof N as P extends "type" ? never : P]: N[P] }
}

which is equivalent to

// type NodeMap = { 
//   Group: { children: Node[]; }; 
//   Obj: { payload: number; };
// }

And then all the operations should be represented in terms of that base type and indexes into that type with a generic key, or generic indexes into a mapped type on that type.

We can recreate Node as such a indexed mapped type (called a distributive object type in microsoft/TypeScript#47109):

type MyNode<K extends keyof NodeMap = keyof NodeMap> =
    { [P in K]: { type: P } & NodeMap[P] }[K]

and then createNode() is written in terms of the generic index K and MyNode<K>:

function createNode<K extends keyof NodeMap>(node: MyNode<K>) {
    const m: { [K in keyof NodeMappings]: (n: MyNode<K>) => NodeMappings[K] } = {
        Group(n) { return new Group(n.children) },
        Obj(n) { return new Obj(n.payload) }
    };
    return m[node.type](node); // okay
}

Note how we had to abandon the control-flow based implementation with switch/case. You can't just check node.type and expect K to be affected, at least until and unless microsoft/TypeScript#33014 is implemented. So instead of doing this, we just make an object whose methods have the same names as node.type and which accept the corresponding node as input. That is, we write an object that maps MyNode<K> to NodeMappings[K].

And this works; the call m[node.type](node) is effectively the same as your switch/case version, but we use property lookups to distinguish between the cases.

And let's make sure it still works for callers:

function test() {
    const n = createNode({ type: 'Obj', payload: 8 });
    //    ^? const n: Obj
}

Looks good. K is inferred as "Obj" and thus the return type is Obj as expected.

Playground link to code