How to make return values conditional on discriminated union argument?-CodePudding

I have a function that takes a discriminated union type as an argument and shall return a different type based on which variant of the argument union was passed. I tried to implement this with conditional types, as recommended in the TypeScript guide. However, TypeScript does not seem to recognize that I am narrowing the original union and will hence not let me assign a return value. Consider the example below:

type Input = { kind: "a" } | { kind: "b" } | { kind: "c" }
type Output<T extends Input> = T extends { kind: "a" }
  ? { a: string }
  : T extends { kind: "b" }
  ? { b: string }
  : { c: string }

export function run<T extends Input>(input: T): Output<T> {
  const kind = input.kind

  switch (kind) {
    case "a":
      return { a: "a" }
    case "b":
      return { b: "b" }
    case "c":
      return { c: "c" }
  }
}

For each of the return statements, TypeScript reports that Type '{ a: string; }' is not assignable to type 'Output<T>'.

How do I fix this error? Are Conditional Types and generics the right tool to determine the output type?

CodePudding user response：

In order to make it work, you need to overload your function:

type Input = { kind: "a" } | { kind: "b" } | { kind: "c" }
type Output<T extends Input> = T extends { kind: "a" }
    ? { a: string }
    : T extends { kind: "b" }
    ? { b: string }
    : { c: string }

/**
 * You need to overload your function,
 * however it is not safe, because you can return {b: 'b'} when kind is 'a'
 */
export function run<T extends Input>(input: T): Output<T>
export function run<T extends Input>(input: T) {
    const kind = input.kind

    switch (kind) {
        case "a":
            return { a: "a" }
        case "b":
            return { b: "b" }
        case "c":
            return { c: "c" }
    }
}

Also, it is not safe, because of this:

type Input = { kind: "a" } | { kind: "b" } | { kind: "c" }
type Output<T extends Input> = T extends { kind: "a" }
    ? { a: string }
    : T extends { kind: "b" }
    ? { b: string }
    : { c: string }

/**
 * You need to overload your function,
 * however it is not safe, because you can return {b: 'b'} when kind is 'a'
 */
export function run<T extends Input>(input: T): Output<T>
export function run<T extends Input>(input: T) {
    const kind = input.kind

    switch (kind) {
        case "a":
            return { b: "a" }      
    }
}

You can also try this:

const STRATEGY = {
    a: { a: 'a' },
    b: { b: "b" },
    c: { c: "c" }
}

const fn = <Key extends keyof typeof STRATEGY>(key: Key): typeof STRATEGY[Key] => STRATEGY[key]

fn('a')