Search code examples
typescriptarrow-functionscartesian-productunion-types

Distribute union of arrow function argument(s) in TypeScript


I'll set out my entire thought process first before getting to the question -- partly because a search engine might pick up on some keywords, and also so that answer givers don't have to elaborate unnecessarily.

In TypeScript, a union type U of two types A and B looks like

type U = A | B;

This means that a value of type U can be of type A or B.

Let's say I want to define an arrow function type F with two parameters (x of type X and y of type Y) that returns a value of type Z. That looks like:

type F = (x: X, y: Y) => Z;   // (1)

So far so good. However, now I want functions of type F to accept either X1 or X2 as the type of their first argument, and Y1 or Y2 as the type of their second argument. From the above, intuition would suggest

type F = (x: X1|X2, y: Y1|Y2) => Z;   // (2)

... but this is wrong! A value of type (x: X1, y: Y1) => Z is not assignable to a value of type F, because the above TypeScript type F stands for functions that have the union type X1|X2 as the type of their first argument, and the union type Y1|Y2 as the type of their second argument. Indeed, (2) is equivalent to:

type Xs = X1 | X2;
type Ys = Y1 | Y2;
type F = (x: Xs, y: Ys) => Z;   // (3), equivalent to (2)

Manually, our problem is solved by instead writing

type F = ( (x: X1, y: Y1) => Z ) | ( (x: X1, y: Y2) => Z ) | ( (x: X2, y: Y1) => Z ) | ( (x: X2, y: Y2) => Z )  // (4)

The union of arguments has been replaced by a union of functions with the former union's types distributed over them. It looks a bit like a Cartesian product of TypeScript types, if you will.

Here's my problem: I only know the union types of the acceptable argument types (Xs and Ys), so I cannot manually write out such a Cartesian product in advance. Going back to my first line of code (the definition of type U), clearly, if only U was given without A and B, you could still use it to mean any one of the types in the union will do. How can I specify any one of the types in the union will do for the type definition of the argument(s) of the arrow function?


Solution

  • Here's a way to do it for functions with two arguments: I've assumed you want the result to be a tagged union, so that you can know what type of arguments to pass to the function at runtime.

    In order to describe the tags for each parameters' union types, it accepts objects like {x1: X1, x2: X2} so that 'x1' and 'x2' become the tags for the X union, and the result uses tags like 'x1-y1'. It would be straightforward to add properties like xTag: Xi, yTag: Yi if you need to preserve the original tags.

    type CreateTaggedUnionOf2Functions<X, Y, T> = {
        [Xi in keyof X]: {
            [Yi in keyof Y]: {
                tag: `${Xi & string}-${Yi & string}`,
                func: (x: X[Xi], y: Y[Yi]) => T
            }
        }[keyof Y]
    }[keyof X]
    
    
    // --- Test ---
    type Test = CreateTaggedUnionOf2Functions<{x1: X1, x2: X2}, {y1: Y1, y2: Y2}, Z>
    /* equivalent to:
    type Test = {tag: "x1-y1", func: (x: X1, y: Y1) => Z}
              | {tag: "x1-y2", func: (x: X1, y: Y2) => Z}
              | {tag: "x2-y1", func: (x: X2, y: Y1) => Z}
              | {tag: "x2-y2", func: (x: X2, y: Y2) => Z}
    */
    

    Playground Link

    I was not satisfied with writing a solution which works specifically and only for functions with two parameters, so here is a generalised version: it is rather complicated, but less so than I thought it might be. Perhaps it can be simplified somehow.

    // --- Implementation ---
    type CreateTaggedUnion<A extends any[]>
        = A extends [infer X]
        ? {
            [Xi in keyof X]: {tag: Xi, arr: [X[Xi]]}
        }[keyof X]
        : A extends [infer X, ...infer Y]
        ? {
            [Xi in keyof X]: CreateTagMap<Y> extends infer U ? {
                [Yi in keyof U]: {
                    tag: `${Xi & string}-${Yi & string}`,
                    arr: U[Yi] extends any[] ? [X[Xi], ...U[Yi]] : never
                }
            }[keyof U] : never
        }[keyof X]
        : never
    
    type CreateTagMap<A extends any[]>
        = CreateTaggedUnion<A> extends infer U
        ? [U] extends [{tag: infer K, arr: any}]
        ? {[Ki in K & string]: (U & {tag: Ki})['arr']}
        : never
        : never
    
    type CreateTaggedUnionOfFunctions<A extends any[], R>
        = CreateTagMap<A> extends infer U
        ? {
            [K in keyof U]: U[K] extends any[] ? {tag: K, func: (...args: U[K]) => R} : never
        }[keyof U]
        : never
    
    
    // --- Tests ---
    interface X1 {x: 1}
    interface X2 {x: 2}
    interface Y1 {y: 1}
    interface Y2 {y: 2}
    interface Z1 {z: 1}
    interface Z2 {z: 2}
    interface O {o: 3}
    
    type TestM = CreateTagMap<[{x1: X1, x2: X2}, {y1: Y1, y2: Y2}, {z1: Z1, z2: Z2}]>
    /** equivalent to TypeScript type:
     * type TestM = {
     *   "x1-y1-z1": [X1, Y1, Z1];
     *   "x1-y1-z2": [X1, Y1, Z2];
     *   "x1-y2-z1": [X1, Y2, Z1];
     *   "x1-y2-z2": [X1, Y2, Z2];
     *   "x2-y1-z1": [X2, Y1, Z1];
     *   "x2-y1-z2": [X2, Y1, Z2];
     *   "x2-y2-z1": [X2, Y2, Z1];
     *   "x2-y2-z2": [X2, Y2, Z2];
     * }
     */
    type TestU = CreateTaggedUnion<[{x1: X1, x2: X2}, {y1: Y1, y2: Y2}, {z1: Z1, z2: Z2}]>
    /** equivalent to TypeScript type:
     * type TestU = {
     *   tag: "x1-y1-z1";
     *   arr: [X1, Y1, Z1];
     * } | {
     *   tag: "x1-y1-z2";
     *   arr: [X1, Y1, Z2];
     * } | {
     *   tag: "x1-y2-z1";
     *   arr: [X1, Y2, Z1];
     * } | {
     *   tag: "x1-y2-z2";
     *   arr: [X1, Y2, Z2];
     * } | {
     *   tag: "x2-y1-z1";
     *   arr: [X2, Y1, Z1];
     * } | ...
     */
    type TestF = CreateTaggedUnionOfFunctions<[{x1: X1, x2: X2}, {y1: Y1, y2: Y2}, {z1: Z1, z2: Z2}], O>
    /** equivalent to TypeScript type:
     * type TestF = {
     *   tag: "x1-y1-z1";
     *   func: (args_0: X1, args_1: Y1, args_2: Z1) => O;
     * } | {
     *   tag: "x1-y1-z2";
     *   func: (args_0: X1, args_1: Y1, args_2: Z2) => O;
     * } | {
     *   tag: "x1-y2-z1";
     *   func: (args_0: X1, args_1: Y2, args_2: Z1) => O;
     * } | {
     *   tag: "x1-y2-z2";
     *   func: (args_0: X1, args_1: Y2, args_2: Z2) => O;
     * } | {
     *   tag: "x2-y1-z1";
     *   func: (args_0: X2, args_1: Y1, args_2: Z1) => O;
     * } | ...
     */
    

    Playground Link