TypeScript の関数で、オブジェクトのメンバのうち引数で指定したものを、そのメンバだけの型で返す

問題のあるコード

const data = {
  hoge: {
    a: 1,
    b: 2,
  },
  fuga: {
    x: 3,
    y: 4,
  },
};

function foo(key: keyof typeof data) {
  return data[key];
}

const d = foo("hoge");
console.log(d.a);

TypeScript: TS Playground - An online editor for exploring TypeScript and JavaScript

関数 foo は、オブジェクト data のメンバ hoge または fuga の内容を、引数 key によって返す関数である。
引数 key の型は、"hoge" | "fuga" となっている。

しかし、この実装では、関数 foo の返り値の型は { a: number; b: number; } | { x: number; y: number; } となってしまい、{ x: number; y: number; } の部分にメンバ a が含まれないため、返り値 d に対するメンバ a へのアクセス d.a はエラーになってしまう。

一方、直感的には、この関数に固定の引数 "hoge" を渡したときの返り値 d は必ず { a: 1, b: 2 } になるため、d.a へのアクセスは許されてほしい。

対処法をAIに聞く

この件について、以下のようにAIに聞いてみた。

以下の TypeScript のコード

const data = Object.freeze({
  hoge: Object.freeze({
    a: 1,
    b: 2,
  }),
  fuga: Object.freeze({
    x: 3,
    y: 4,
  }),
});

function foo(key: keyof typeof data) {
  return data[key];
}

現状では foo の返り値の型は hoge の型と fuga の型のORになっているが、これをkeyの値に応じて一方だけが返るようにしたい。どうすればいい？

Perplexity AI: (質問文省略) (魚拓)

前述の例とはコードが異なるが、これはAIに聞いてみた後、よりシンプルな例でこの記事を書くことにしたためである。

得られた解決策

AIの回答により、今回の例では以下のように書き換えればいいことがわかった。

const data = {
  hoge: {
    a: 1,
    b: 2,
  },
  fuga: {
    x: 3,
    y: 4,
  },
};

function foo<K extends keyof typeof data>(key: K): (typeof data)[K] {
  return data[key];
}

const d = foo("hoge");
console.log(d.a);

TypeScript: TS Playground - An online editor for exploring TypeScript and JavaScript

d の型は { a: number; b: number; } となり、d.a に問題なくアクセスできる。

この対処法は、Generic Functions の Constraints を用いているようである。
AIの回答では「Conditional Types を使うとよい」と主張しているが、この要素はなさそうである。

さらに、この例では、返り値の型は明示しなくても動くようである。

const data = {
  hoge: {
    a: 1,
    b: 2,
  },
  fuga: {
    x: 3,
    y: 4,
  },
};

function foo<K extends keyof typeof data>(key: K) {
  return data[key];
}

const d = foo("hoge");
console.log(d.a);

TypeScript: TS Playground - An online editor for exploring TypeScript and JavaScript

まとめ

今回の例では、AIの回答に含まれる解説には若干の誤りがみられたものの、AIが提示してくれたコードの修正方法は的確であり、問題の解決に繋がった。