I have an object of type Foo, which has a member bar: Bar. Bar is a tagged union, which I have defined as

type ABar = {name: 'A', aData: string};
type BBar = {name: 'B', bData: string};
type Bar = ABar | BBar;
type BarName = Bar['name']; // this is equivalent to 'A'|'B'

Different parts of my code want to process foo.bar only if it has a particular type; i.e. some parts only care if the bar is an ABar, others only care if the bar is a BBar. I'd like to define a function that extracts the Bar from a Foo in a type safe way. What I have so far is a function:

function getBar(foo: Foo, name: BarName): Bar | null {
    return foo.bar.name === name ? foo.bar : null;
}

This works, but the caller has to explicitly assert the correct type, i.e.

let aBar = getBar(foo, 'A') as ABar | null;

Is there any way to get the TS compiler to infer that if the given name is 'A', then the result must indeed have type ABar (or null)?

Here is the same example in TS Playground.

CodePudding user response：

You could make the function generic and use the Extract utility type to extract a union member based on a given name.

function getBar<T extends BarName>(foo: Foo, name: T) {
    return (
      foo.bar.name === name ? foo.bar : null
    ) as Extract<Bar, { name: T }> | null;
}

let aBar = getBar(foo, 'A');
// let aBar: ABar | null

let bBar = getBar(foo, 'B');
// let bBar: BBar | null

Playground

CodePudding user response：

First, it's worth noting that TypeScript will infer the correct type of Bar if you check the name property as part of an if statement:

if (bar.name === "A") {
  bar.aData; // No error, bar has type "ABar" here.
}

This might be good enough depending on how your code is structured. If you still want to use the getBar function though, you might be able to use one of these approaches:

Option one, you could use another type that associates each of the Bar types to it's name. Eg:

type BarMap = {
  A: ABar;
  B: BBar;
}

This makes looking up the correct associated type, given the name, very straightforward if you use a generic function:

function getBar<T extends BarName>(foo: Foo, name: T): BarMap[T] {
  // @ts-ignore: SMH TypeScript, this works fine.
  return foo.bar.name === name ? foo.bar : null;
}

Unfortunately this creates a nasty error message on the return line, but ultimately I would just ignore that error since this will work fine at runtime and it achieves the correct inference:

let aBar = getBar(foo, 'A'); // Type: ABar | null
let bBar = getBar(foo, 'B'); // Type: BBar | null

If you try and pass it a name that is not part of the BarName type, it will produce an error:

let cBar = getBar(foo, 'C') // Error: Argument of type '"C"' is not assignable to parameter of type '"A" | "B"'

Option two is to use an overload signature on your function, which will allow you to specify the exact return type for each parameter combination. Eg:

function getBar(foo: Foo, name: 'A'): ABar | null;
function getBar(foo: Foo, name: 'B'): BBar | null;
function getBar(foo: Foo, name: BarName): Bar | null {
  return foo.bar.name === name ? foo.bar : null;
}

This also achieves the same inference and will also produce an error if you pass a non-valid name. If you wanted to, you could also add another signature to the overload to set the return type to a bare null for unknown keys.

Depending on your use case, you may find one of these solutions to be more maintainable than the other.

Playground Link