Search code examples
typescriptdiscriminated-union

"Switch" function over discriminated union members in Typescript


Suppose I have a discriminated union in Typescript:

type MyAction = { type: "redirect", url: string} | {type: "returnValue", value: number}

I want to write a "switch expression" which takes a member of the discriminated union and a mapping of types to processor functions, and returns the result of running the correct processor function:

let actionInstance = {type: "returnValue", value: 123} as MyAction
console.log(switchExpression(actionInstance, {
  redirect: (action) => action.url + "(url)",
  returnValue: (action) => `${action.value}` + "(value)"}
))

I want the processors to be typed - that is, I want Typescript to error if not all types are represented in the switch, if I provide a key for the switch which is not a valid type, or if I assume the wrong type in a processor function (like I try to access action.value in the value for redirect).

This is what I've got so far (playground link), based on a few other answers on StackOverflow:

type NarrowAction<T, N> = T extends { type: N } ? T : never;

type StuctureRecord<FunctionReturnType> =
    { [K in MyAction["type"]]: (value: NarrowAction<MyAction, K>) => FunctionReturnType }

function switchExpression<FunctionReturnType>(value: MyAction, cases: StuctureRecord<FunctionReturnType>): FunctionReturnType {
  let f = cases[value.type] as (value: any) => FunctionReturnType
  return f(value)
}

This works correctly, but it's specific to one discriminated union (MyAction).

Is there a way to adapt this code to work with any discriminated union where all the members have a type field, instead of having it only work for MyAction? I tried using a type like DUType extends { type: string }, but that results in a bunch more errors (playground link)

Note: I know I can use the StructuredRecord directly or I could just use a regular switch statement or if/else statements with an assertUnreachable function (StackOverflow post), but I'm interested in whether a solution in the above form exist due to academic curiosity, and that fact that it would be an ergonomic util function to have to ensure type safety without having to type all the type annotations directly.


Solution

  • If you have a discriminated union where the type property is a discriminant, you can expand StructureRecord to:

    type StuctureRecord<A extends { type: string }, R> =
      { [T in A as T["type"]]: (value: T) => R }
    

    where R corresponds to your FunctionReturnType type parameter (I like to follow the naming convention so you can easily tell a type parameter apart from a specific type), and A is the discriminated union type.

    This is similar to yours except I'm using key remapping to iterate over the entire discriminated union type A instead of over the type property of A. By doing T in A I'm able to get each member of the union type in term, and just use it as the function parameter type. Otherwise you have to do K in A["type"] and then use K to recreate T from A with something like the Extract<T, U> utility type, like you're doing with NarrowAction. This change is just for elegance/convenience and doesn't really affect any behavior we care about.


    At this point we could just add another generic type parameter A to switchExpression corresponding to the discriminated union:

    const switchExpression =
      <A extends { type: string }, R>(value: A, cases: StuctureRecord<A, R>): R => {
        const type: A["type"] = value.type;
        const f = cases[type] as (value: any) => R;
        return f(value);
      }
    

    Note that I have moved type to its own variable and annotated it as type A["type"]. If you just write value.type the compiler will eagerly resolve that to string because A extends {type: string}. But we don't want that, since cases will be known to have keys in A["type"] but not just string. This saves us from needing another type assertion to write cases[type as A["type"]].

    But since we need that cases[type] as (value: any) => R in there, it doesn't matter that much. There's a fundamental limitation of TypeScript where the compiler just can't do the higher order reasoning to see that for each member T of the A union, StructureRecord<T, R>[T["type"]] will be of type (value: T) => R and thus it will always be safe to call cases[value.type](value). The correlation between cases[value.type] and value over the different possible narrowings of A is lost. See microsoft/TypeScript#30581 for more information. Thus, inside the implementation of switchExpression you should just use as many type assertions as you need to avoid errors, triple check that your assertions are reasonable and that you're not allowing unsafe things, and move on.


    So that's great. Unfortunately, when you call switchExpression(), you see a problem with type inference:

    switchExpression(actionInstance, {
      redirect: (action) => action.url + "(url)",
      // ------> ~~~~~~ // error, action is implicitly any
      returnValue: (action) => `${action.value}` + "(value)"
      // ---------> ~~~~~~ // error, action is implicitly any    
    }).toUpperCase(); // <-- error, object is of type unknown
    

    Here the compiler does infer A to be MyAction, but the callback parameters of action are implicitly any and not the expected members of the MyAction union. This is another limitation of TypeScript. In this case we need the compiler to infer both the type parameters A and R as well as the types of the two callback parameters named action based on context. The compiler tries to do this all at once and fails to get action right because it depends on A which it doesn't know yet. And R also fails to be inferred, falls back to unknown, and then the compiler won't let you use the return value easily.

    There are various issues in GitHub around this problem; see microsoft/TypeScript#25092 for one of them. Whenever you want both generic type inference and contextual callback parameter type inference at once, you run the risk of this kind of issue.

    You could cut the knot by manually specifying both A and R:

    switchExpression<MyAction, string>(actionInstance, {
      redirect: (action) => action.url + "(url)",
      returnValue: (action) => `${action.value}` + "(value)"
    }).toUpperCase(); // okay
    

    but that's a bit sad.


    Instead we could work around the problem by making switchExpression into a curried function which takes value and returns a function that takes cases. The original function will be able to infer A properly, after which the compiler only needs to infer R and action, because A is already set in stone:

    const switchExpression =
      <A extends { type: string }>(value: A) =>
        <R,>(cases: StuctureRecord<A, R>): R => {
          const type: A["type"] = value.type;
          const f = cases[type] as (value: any) => R;
          return f(value);
        }
    

    That's not much of a change, really. And you call it like this:

    switchExpression(actionInstance)({
      redirect: (action) => action.url + "(url)",
      returnValue: (action) => `${action.value}` + "(value)"
    }).toUpperCase(); // okay
    

    If you squint you can think of )( as being like , when you call switchExpression. But you can see that type inference has occurred exactly as we want. A is inferred as MyAction in the first call, and so can be used to contextually type action in the second call, which allows the compiler to infer R in turn.

    Personally I'd rather use currying and get nice type inference than a single function call and have to manually specify types, but it's up to you what makes sense for your use cases.

    Playground link to code