【TypeScript】引数の型によって返り値の型を変える関数

追記: 関数のオーバーロードを使うほうがいい

こういう関数を作りたいとする。

type Func =
  | ((arg: "STRING") => string)
  | ((arg: "NUMBER") => number)

• 引数に "STRING" を与えると string 型が返ってくる
• 引数に "NUMBER" を与えると number 型が返ってくる

このような関数の型を定義するには Conditional Type を使えばよい。

// 引数の Union 型
type Arg = "STRING" | "NUMBER"

// T が "STRING" なら string 型、"NUMBER" なら number 型
type Return<T> = T extends "STRING"
  ? string
  : T extends "NUMBER"
  ? number
  : never

const func = <T extends Arg>(a: T): Return<T> => {
  switch (a) {
    case "STRING":
      return "a" as Return<T>
    case "NUMBER":
      return 1 as Return<T>
  }
}

// 次のように型推論される: const func: <"NUMBER">(a: "NUMBER") => number
func("NUMBER")

// 次のように型推論される: const func: <"STRING">(a: "STRING") => string
func("STRING")