如何在typescript中检查function中的类型联合体的类型

Question

Say I have this code:

interface A {
 type: 'a',
 value : number
}

interface B {
  type: 'b',
  value: string
}

type AB = A | B;

const db: Record<string, AB> = {
  param1: { type: 'a', value: 1234 }
};

I have code like this (over-simplified) that works well

function processParameter(parameter: string, expectedType: 'a' | 'b') {
  if (expectedType === 'a') {
    const result = db[parameter];
    if(!result) { throw new Error('Not found'); }
    if(result.type !== 'a') { throw new Error('Unexpected type'); }
    processA(result);
  }
  if (expectedType === 'b') {
    const result = db[parameter];
    if(!result) { throw new Error('Not found'); }
    if(result.type !== 'b') { throw new Error('Unexpected type'); }
    processB(result);
  }
}

function processA(value: A) {
  // ...
}

function processB(value: B) {
  // ...
}

I would like to refactor the code like this:

function processParameter(parameter: string, expectedType: 'a' | 'b') {
  if (expectedType === 'a') {
    const result = getAndCheck(parameter, expectedType);
    processA(result);
  }
  if (expectedType === 'b') {
    const result = getAndCheck(parameter, expectedType);
    processB(result);
  }
}

// I don't want to write all those lines here, is there a way ?
function getAndCheck(parameter: string, expectedType: 'a'): A;
function getAndCheck(parameter: string, expectedType: 'b'): B;
function getAndCheck(parameter: string, expectedType: 'a' | 'b'): AB {
  const result = db[parameter];
  if(!result) { throw new Error('Not found'); }
  if(result.type !== expectedType) { throw new Error('Unexpected type'); }
  return result;
}

Is there a way to simplify that? Using generic? using inference? am I going the wrong way?

Answer 1

You can use a generic function to establish a relation between the expectedType parameter and the return type:

function getAndCheck<K extends AB['type']>(parameter: string, expectedType: K): Extract<AB, { type: K }> {
    const result = db[parameter] as AB;
    if (!result) { throw new Error('Not found'); }
    if (result.type !== expectedType) { throw new Error('Unexpected type'); }
    return result as Extract<AB, { type: K }>;
}

Playground Link

Here we use a type parameter K to capture the actual type that is passed in and we ensure that only a type in the union AB can be passed in (we use an index access type to get the union of all possible types). We then use Extract to get the appropriate element from the union. (You can read a bit more abotut filtering unions here )