return type for pattern matching function in typescript

Question

I am trying to create a pattern-matching function for typescript that works on a discriminated union.

For example:

export type WatcherEvent =
  | { code: "START" }
  | {
      code: "BUNDLE_END";
      duration: number;
      result: "good" | "bad";
    }
  | { code: "ERROR"; error: Error };

I want to be able to type a match function that looks like this:

match("code")({
    START: () => ({ type: "START" } as const),
    ERROR: ({ error }) => ({ type: "ERROR", error }),
    BUNDLE_END: ({ duration, result }) => ({
      type: "UPDATE",
      duration,
      result
    })
})({ code: "ERROR", error: new Error("foo") });

So far I have this

export type NonTagType<A, K extends keyof A, Type extends string> = Omit<
  Extract<A, { [k in K]: Type }>,
  K
>;

type Matcher<Tag extends string, A extends { [k in Tag]: string }> = {
  [K in A[Tag]]: (v: NonTagType<A, Tag, K>) => unknown;
};

export const match = <Tag extends string>(tag: Tag) => <
  A extends { [k in Tag]: string }
>(
  matcher: Matcher<Tag, A>
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
) => <R extends any>(v: A): R =>
  (matcher as any)[v[tag]](v);

But I do not know how to type the return type of each case

At the moment each case is typing the parameters correctly but the return type is unknown so if we take this case

ERROR: ({ error }) => ({ type: "ERROR", error }), // return type is not inferred presently

then the return type of each case like function is unknown as is the return type of match function itself:

Here is a codesandbox.

Answer 1

The way I see it there are two approaches that you can take with this.

1. The input type is known beforehand

If you want to enforce that the initialisation of the final function takes a particular type then that type must be known beforehand:

// Other types omitted for clarity:
const match = <T>(tag) => (transforms) => (source) => ...

In this example you specify T at the time of the first call, with the following type constraints as a consequence:

1. tag must be a key of T
2. transforms must be an object with keys for all values of T[typeof tag]
3. source must be of type T

In other words, the type that substitutes T determines the values that tag, transforms and source can have. This seems the most straightforward and understandable to me, and I'm going to try to give an example implementation for this. But before I do, there's also approach 2:

2. the input type is inferred from the last call

If you want to have more flexibility in the type for source based on the values for tag and transforms, then the type can be given at, or inferred from, the last call:

const match = (tag) => (transforms) => <T>(source) => ...

In this example T is instantiated at the time of the last call, with the following type constraints as a consequence:

1. source must have a key tag
2. typeof source[tag] must be a union of at most all transforms keys, i.e. keyof typeof transforms. In other words, (typeof source[tag]) extends (keyof typeof transforms) must always be true for a given source.

This way, you are not constrained to a specific substitution of T, but T might ultimately be any type that satisfies the above constraints. A major downside to this approach is that there will be little type checking for the transforms, since that can have any shape. Compatibility between tag, transforms and source can only be checked after the last call, which makes things a lot harder to understand, and any type-checking errors will probably be rather cryptic. Therefore I'm going for the first approach below (also, this one is pretty tough to wrap my head around ;)

---

Because we specify the type in advance, that's going to be a type slot in the first function. For compatibility with the further parts of the function it must extend Record<string, any>:

const match = <T extends Record<string, any>>(tag: keyof T) => ...

The way we would call this for your example is:

const result = match<WatcherEvent>('code') (...) (...)

We are going to need the type of tag for further building the function, but to parameterise that, e.g. with K would result in an awkward API where you have to write the key literally twice:

const match = <T extends Record<string, any>, K extends keyof T>(tag: K)
const result = match<WatcherEvent, 'code'>('code') (...) (...)

So instead I'm going for a compromise where I'll write typeof tag instead of K further down the line.

Next up is the function that takes the transforms, let's use the type parameter U to hold its type:

const match = <T extends Record<string, any>>(tag: keyof T) => (
    <U extends ?>(transforms: U) => ...
)

The type constraint for U is where it gets tricky. So U must be an object with one key for each value of T[typeof tag], each key holding a function that transforms a WatcherEvent to anything you like (any). But not just any WatcherEvent, specifically the one that has the respective key as its value for code. To type this we'll need a helper type that narrows down the WatcherEvent union to one single member. Generalising this behaviour I came up with the following:

// If T extends an object of shape { K: V }, add it to the output:
type Matching<T, K extends keyof T, V> = T extends Record<K, V> ? T : never

// So that Matching<WatcherEvent, 'code', 'ERROR'> == { code: "ERROR"; error: Error }

With this helper we can write the second function, and fill in the type constraint for U as follows:

const match = <T extends Record<string, any>>(tag: keyof T) => (
    <U extends { [V in T[typeof tag]]: (input: Matching<T, typeof tag, V>) => any }>(transforms: U) => ...
)

This constraint will make sure that all function input signatures in transforms fit the inferred member of the T union (or WatcherEvent in your example).

Note that the return type any here does not loosen the return type ultimately (because we can infer that later on). It simply means that you're free to return anything you want from functions in transforms.

Now we've come to the last function -- the one that takes the final source, and its input signature is pretty straightforward; S must extend T, where T was WatcherEvent in your example, and S is going to be the exact const shape of the given object. The return type uses the ReturnType helper of typescript's standard library to infer the return type of the matching function. The actual function implementation is equivalent to your own example:

const match = <T extends Record<string, any>>(tag: keyof T) => (
    <U extends { [V in T[typeof tag]]: (input: Matching<T, typeof tag, V>) => any }>(transforms: U) => (
        <S extends T>(source: S): ReturnType<U[S[typeof tag]]> => (
            transforms[source[tag]](source)
        )
    )
)

That should be it! Now we could call match (...) (...) to obtain a function f that we can test against different inputs:

// Disobeying some common style rules for clarity here ;)

const f = match<WatcherEvent>("code") ({
    START       : ()                     => ({ type: "START" }),
    ERROR       : ({ error })            => ({ type: "ERROR", error }),
    BUNDLE_END  : ({ duration, result }) => ({ type: "UPDATE", duration, result }),
})

And giving this a try with the different WatcherEvent members gives the following result:

const x = f({ code: 'START' })                                     // { type: string; }
const y = f({ code: 'BUNDLE_END', duration: 100, result: 'good' }) // { type: string; duration: number; result: "good" | "bad"; }
const z = f({ code: "ERROR", error: new Error("foo") })            // { type: string; error: Error; }

Note that when you give f a WatcherEvent (union type) instead of a literal value, the returned type will also be the union of all return values in transforms, which seems like the proper behaviour to me:

const input: WatcherEvent = { code: 'START' }
const output = f(input)

// typeof output == { type: string; }
//                | { type: string; duration: number; result: "good" | "bad"; }
//                | { type: string; error: Error; }

Lastly, if you need specific string literals in the return types instead of the generic string type, you can do that by just altering the functions that you define as transforms. For example you could define an additional union type, or use 'as const' annotations in the function implementations.

Here's a TSPlayground link, I hope this is what you're looking for!