Narrowing return value of function based on argument

user_578706 0
Viewed 60

Consider the following code

type Foo = {
  kind: "foo";
  foo: number;
};

type Bar = {
  kind: "bar";
  bar: number;
};

type FooBar = Foo | Bar;

function foobar(arg: FooBar, kind: "foo" | "bar"): FooBar | undefined {
  if (arg.kind === kind) {
    return arg;
  }
  return undefined;
}

declare const a: FooBar;

if (a.kind === "foo") {
  a.foo;  // no error, type is correctly narrowed
}

foobar(a, "foo").foo; // error: Property 'foo' does not exist on type 'Bar'

Is it possible to write type signature for foobar function, so return value is narrowed according to the kind argument?

typescript
2 Answers

We can extract the kinds by using generic parameters to model the types that we want to use as 'variables' across our output types.

In this case I use ExtractKind<T> to extract the unions of different kinds of T.

Then we can use RefineKind<T, K> - where T is the type with a kind and K is the kind to refine to.

I've wrapped the output result with an undefined check - if you want I can provide an example on how to exhaustively map all types to remove the need for this check but this seems outside the scope of the original question.

Let me know if you need any further explanation on these concepts!

Full example:


type Foo = {
  kind: "foo";
  foo: number;
};

type Bar = {
  kind: "bar";
  bar: number;
};

type FooBar = Foo | Bar;

type ExtractKind<T> = T extends { kind: infer K } ? K : never;
type RefineKind<T, K> = T extends { kind: K } ? T : never;

function foobar<K extends ExtractKind<FooBar>>(arg: FooBar, kind: K) {
  if (arg.kind === kind) {
    return arg as RefineKind<FooBar, K>;
  }
  return undefined;
}

const foo = foobar(a, "foo");

if (foo !== undefined) {
  foo.foo; // foo -> number
}

Typescript playground link here

user_23104500

I think what you're really trying to do here is implement a user-defined type guard. Essentially you write a boolean function with the return type arg is SomeType and that tells typescript that if the function is true than arg must be of type Sometype. And if it's false then that means that arg cannot be of type Sometype.

Here is what that function looks like in your case:

function isKind<T extends Kinds>(arg: FooBar, kind: T): arg is Extract<FooBar, {kind: T}> {
    return arg.kind === kind;
}

When you use this function inside an if statement, you get meaningful information about the type:

function testUnknown( either: FooBar ) {
    if ( isKind(either, "bar") ) {
        // we now know that either is Bar
        const barVal = either.bar; // ok
        const fooVal = either.foo; // error
    }
    else {
        // we now know that either is Foo
        const barVal = either.bar; // error
        const fooVal = either.foo; // ok
    }
}

I personally would prefer to write isKind as a double arrow function so that you could more easily call isKind("foo") inside an array.map or other callback, but that's just down to code style and personal preference.

const isKind = <T extends Kinds>(kind: T) => 
    (arg: FooBar): arg is Extract<FooBar, {kind: T}> => {
        return arg.kind === kind;
    }

Your original foobar function can use your isKind type guard internally. This is similar to @mbdavis's answer but you don't need to assert the type of arg using "as" because typescript already understands that the type has been refined.

function foobar<T extends Kinds>(arg: FooBar, kind: T): Extract<FooBar, {kind: T}> | undefined {
  return isKind(arg, kind) ? arg : undefined;
}

But I suspect that you don't really even need this foobar function and can just call isKind directly in your code.

Typescript Playground Link

user_104315740
Related