Automatically infer return type based on arguments for discriminated union in TypeScript

Issue

I am trying to filter an array and infer the return type automatically.

enum Category {
  Fruit,
  Animal,
  Drink,
}

interface IApple {
  category: Category.Fruit
  taste: string
}

interface ICat {
  category: Category.Animal
  name: string
}

interface ICocktail {
  category: Category.Drink
  price: number
}

type IItem = IApple | ICat | ICocktail

const items: IItem[] = [
  { category: Category.Drink, price: 30 },
  { category: Category.Animal, name: 'Fluffy' },
  { category: Category.Fruit, taste: 'sour' },
]

So now I want to filter the items, something like:

// return type is IItem[], but I want it to be IFruit[]
items.filter(x => x.category === Category.Fruit)

I understand that Array#filter is too generic to do that, so I'm trying to wrap it in a custom function:

const myFilter = (input, type) => {
  return input.filter(x => x.category === type)
}

So, all I need is add types and it's good to go. Let's try:

First idea is to add return conditional types:

const myFilter = <X extends IItem, T extends X['category']>(
  input: X[],
  type: T
): T extends Category.Fruit ? IApple[] : T extends Category.Drink ? ICocktail[] : ICat[] => {
  // TS error here
  return input.filter((x) => x.category === type)
}

While the return type of myFilter is now indeed works well, there are 2 problems:

• input.filter((x) => x.category === type) is highlighted as a error: Type 'X[]' is not assignable to type 'T extends Category.Fruit ? IApple[] : T extends Category.Drink ? ICocktail[] : ICat[]'
• I manually specified all possible cases, basically doing the work compiler should do. It's easy to do when I have only 3 type intersections, but when there are 20... not so easy.

Second idea was to add some sort of a constraint, like so:

const myFilter = <X extends IItem, T extends X['category'], R extends ...>(input: X[], type: T): X[] => {
  return input.filter(x => x.category === type)
}

but what R extends? I don't.

The third idea is to use overloading, however, it's not a good idea either as it'll require specifying all types manually, just like in idea #1.

Is it possible in modern TS to solve this problem by using compiler only?

Solution

The issue is not with Array.prototype.filter(), whose typings in the standard TS library actually does have a call signature that can be used to narrow the type of the returned array based on the callback:

interface Array<T> {
  filter<S extends T>(
    predicate: (value: T, index: number, array: T[]) => value is S, 
    thisArg?: any
  ): S[];
}

The issue is that this call signature requires the callback to be a user-defined type guard function, and currently such type guard function signatures are not inferred automatically (see microsoft/TypeScript#16069, the open feature request to support this, for more info). So you'll have to annotate the callback yourself.

And in order to do that generically, you do probably want conditional types; specifically I'd suggest using the Extract<T, U> utility type to express "the member(s) of the T union assignable to the type U":

const isItemOfCategory =
  <V extends IItem['category']>(v: V) =>
    (i: IItem): i is Extract<IItem, { category: V }> =>
      i.category === v;

Here, isItemOfCategory is a curried function that takes a value v of type V assignable to IItem['category'] (that is, one of the Category enum values) and returns a callback function that takes an IItem i and returns a boolean whose value the compiler can use to determine if i is or is not an Extract<IItem, { category: V }>... which is "the member of the IItem union whose category property is of type V". Let's see it in action:

console.log(items.filter(isItemOfCategory(Category.Fruit)).map(x => x.taste)); // ["sour"]
console.log(items.filter(isItemOfCategory(Category.Drink)).map(x => x.price)); // [30]
console.log(items.filter(isItemOfCategory(Category.Animal)).map(x => x.name)); // ["Fluffy"]

Looks good. I don't see the need to try to refactor further into a different type signature for filter(), since the existing one works how you want.

Playground link to code

Answered By - jcalz