Is there any way to have TypeScript infer the type of a function's return value based on the call arguments?

Question

I'm setting up a function that will retrieve data from my database. It takes a key string as argument, and it returns a different data type based on the key.

My goal is to have TypeScript infer the type that will be returned based on the key that's passed as argument.

This is what I've tried:

interface Fruits { names: string[] }
interface Employees { ages: number[] }

function getData(key: "fruits" | "employees") {
  if (key === "fruits") {
    // fetch data for /fruits and return it, specifying its type
    return { names: ["apple", "orange", "kiwi"] } as Fruits;
  }

  // fetch data for /employees and return it, specifying its type
  return { ages: [30, 50, 19 ] } as Employees;
}

const fruits = getData("fruits"); 
// ...should infer that the returned value will be of type 'Fruits'
        
const fruitNames = fruits.names; 
// ...fails to infer type. Thinks fruits is of type (Fruits | Employees)

The last line gives a warning because TypeScript could not infer that using the key "fruits" when calling getData will always return a value of type Fruits (which, in turn, has the property "names").

Of course, I could make the error go away by also typecasting my call to getData like this:

const fruits = getData("fruits") as Fruits;

But, in my opinion, that defeats the purpose of using TypeScript in the first place since once again my code becomes very error-prone. If, for example, I make the wrong typecasting, code that's technically incorrect won't give any warning:

const fruits = getData("fruits") as Employees;
const nonexistent = fruits.ages; 
// TS won't complain about any accesses to the non-existent property 'ages' 
// because I made the wrong typecasting

My goal is that whoever uses the getData function knows, in advance, thanks to TypeScript's IntelliSense, the type of data they will get back as soon as they provide the 'key' string to query data with. It should not be necessary to manually typecast the function's return value type when calling the function.

Inside of getData, any typecasting, use of generics, or any other typing is fine. But, outside of it, giving the function a valid key argument should be enough to have TypeScript infer the correct return type.

Is there any way to achieve this?

Answer 1

Another approach is to declare a function overload for each of the keys with a return value of the appropriate type. Your implementation would type key with all its possible values and have a return type that is a union of all the possible returns. That would look like:

interface Fruits { names: string[] }
interface Employees { ages: number[] }

function getData(key: "fruits"): Fruits;
function getData(key: "employees"): Employees;

function getData(key: "fruits" | "employees"): Fruits | Employees {
  if (key === "fruits") {
    // fetch data for /fruits and return it
    return { names: ["apple", "orange", "kiwi"] };
  }

  // fetch data for /employees and return it
  return { ages: [30, 50, 19 ] };
}

Answer 2

There is actually several ways, here is one way

interface Fruits {
    names: string[];
}
interface Employees {
    ages: number[];
}

interface DataMap {
    fruits: Fruits;
    employees: Employees;
}

function getData<T extends keyof DataMap>(key: T): DataMap[T] {
if (key === "fruits") {
    // fetch data for /fruits and return it, specifying its type
    return { names: ["apple", "orange", "kiwi"] } as DataMap[T];
}

// fetch data for /employees and return it, specifying its type
    return { ages: [30, 50, 19] } as DataMap[T];
}

const fruits = getData("fruits"); // knows it's a Fruits
const fruitNames = fruits.names; // works

We make an auxilliary type which serves to map strings to types, and use this generic to capture T so we can use it in the return type.