How to leverage discriminated union to infer return type of a function

Question

Given the following types, interfaces, and getData function below I'm trying to find a way to leverage discriminated unions so that the TS compiler can narrow the return type of getData(source: DOSources) to the associated DOTypes

// Expected behavior
const result = getData("dataObjectA");

// result.data should be a string but in this case the TS compiler will complain 
// that data does not have the toLowerCase() function
result.data.toLowerCase();

Example Code

interface DataObjectA {
  source: "dataObjectA";
  data: string;
}

interface DataObjectB {
  source: "dataObjectB";
  data: number;
}


type DOTypes = DataObjectA | DataObjectB
type DOSources = DOTypes["source"];

async function getData(source: DOSources) {
  const response = await fetch(`https://some-random-endpoint/`, {
    method: "GET",
    headers: {
      "Content-Type": "application/json",
    },
  });

  switch (source) {
    case "dataObjectA":
      return await response.json() as DataObjectA;
    case "dataObjectB":
      return await response.json() as DataObjectB;
  }
}

Answer 1

You can indeed get the compiler to compute the desired return type of getData() as a function of the DOTypes discriminated union and the type of the source parameter. You can make getData() a generic function whose type parameter K extends DOSources is the type of the source parameter. For example:

async function getData<K extends DOSources>(source: K) {
  const response = await fetch(`https://some-random-endpoint/`, {
    method: "GET",
    headers: {
      "Content-Type": "application/json",
    },
  });

  return await response.json() as Extract<DOTypes, { source: K }>
}

To find the member of the DOTypes discriminated union associated with K, we can use the Extract utility type. Extract<DOTypes, {source: K}> selects from DOTypes all union members whose source property is of a type assignable to K.

Note that we have to assert that the function returns a value of (a Promise corresponding to) this type; the compiler is unable to verify that.

---

Let's test it:

const resultA = await getData("dataObjectA"); // const result: DataObjectA
resultA.data.toLowerCase();

const resultB = await getData("dataObjectB"); // const result: DataObjectB
resultB.data.toFixed();

Looks good. Each result is narrowed to the expected type. You'll only get a union out of getData() if you put a union in:

const resultAOrB = await getData(Math.random() < 0.5 ? "dataObjectA" : "dataObjectB");
// const resultAOrB: DataObjectA | DataObjectB

Playground link to code