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

Answer 2

One option is to use function overloads to specify the different potential return values for different input types. Here is a sandbox link for the following 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: DataObjectA["source"]): Promise<DataObjectA>;
async function getData(source: DataObjectB["source"]): Promise<DataObjectB>;
async function getData(source: DOSources) {
  const response = await fetch(`https://some-random-endpoint/`, {
    method: "GET",
    headers: {
      "Content-Type": "application/json",
    },
  });

  return await response.json();
}

async function test() {
    // string
    const result = await getData("dataObjectA");
    result.data.toLowerCase();
    // number
    const result2 = await getData("dataObjectB");
    result2.data.toFixed(3);
}

That being said, if you're not actually using the source parameter, you could just explicitly pass the type to determine the output rather than passing a variable. Again, a playground link for this option.

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

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

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

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

  return await response.json();
}

async function test() {
  // string
  const result = await getData<"dataObjectA">();
  result.data.toLowerCase();
  // number
  const result2 = await getData<"dataObjectB">();
  result2.data.toFixed(3);
}