Make typescript function return type conditional on input-CodePudding

In the bellow code b will have the type number | string even though b would always be assigned a string if we use Format.hex as an input to the function. Is there a way to do some smart typescript magic e.g. with generics such that b is assigned the correct type based on the function input?

enum Format{
  hex,
  integer
}

function fun (a: Format){
  const k = Math.floor(Math.random()*100)
  if (a === Format.hex) {
    return k.toString(16)
  }
  else { 
    return k
  }
}

const b = fun(Format.hex)

CodePudding user response：

You can make the use of the overload function signature, like so:

enum Format {
  hex,
  integer
}

// Overload signature without a body.
// Here you're saying that if the function gets called with Format.hex it will return string 
function fun(a: Format.hex): string;
function fun(a: Format.integer): number;
function fun(a: Format) {
  const k = Math.floor(Math.random() * 100);

  if (a === Format.hex) {
    return k.toString(16);
  }

  return k;
}

const b = fun(Format.hex); // b is now of type 'string'