Constrain return type from argument function

Question

I have a function taking a function as an argument. The argument function returns one of a few enums. How do I infer the type of the wrapper function by using the return type of the argument function.

I think much easier explained as an example:

const foo = { a: 1, b: '2' }

function h(k: (() => 'a' | 'b')) {
  return foo[k()]
}

const d = h(() => 'a') 
// expected: d should be of type number
//   actual: d is of type number | string

playground

Answer 1

You can define returned key type as generic parameter (it will be inferred as a single key and not union of all possible keys):

function h<K extends keyof typeof foo>(k: (() => K)) {
  return foo[k()]
}

const d = h(() => 'a') // now number

Playground