I'm trying to write the typescript signature of a generic "groupBy" function that would "spread" a discriminated type union array into a record where each field of the record is a possible discriminator value and points to an array of objects of a concrete type from the union.

Example:

interface Dog {
  type: 'dog'
  dogMetadata: {}
}

interface Cat {
  type: 'cat'
  catMetadata: {}
}

type Animal = Dog | Cat

const animals: Animal[] = [{ type: 'dog', dogMetadata: {} }, { type: 'cat', catMetadata: {} }]

Each interface has a common discriminator property and no other common properties.

Here is simple "groupBy" signature that does not spread the type union values, forcing me to downcast the values of the record:

function groupBy<T, K extends string>(arr: T[], keyExtractor: (element: T) => K): Record<K, T[]>

const animalsByType: Record<'dog' | 'cat', Animal[]> = groupBy(animals, it => it.type)
const dogs: Dog[] = animalsByType['dog'] as Dog[]  // Must downcast Animal[] to Dog[]

How can I make a "groupBy" that is aware of the concrete types of the discriminated union type? I would want something like this:

const animalsByType: { dog: Dog[], cat: Cat[] } = groupBy(animals, it => it.type)
const dogs: Dog[] = animalsByType['dog']  // animalsByType.dog is known to be Dog[] by typescript

The implementation is easy, having trouble with the Typescript part :) I'm looking for a generic solution that doesn't make assumptions, like the name of the discriminator property or the amount of types in the type union.

Follow-up question

Would it be possible to make the same signature work when the union is nested inside another class?

interface Holder<T> {
  data: T
}

const animalHolders: Holder<Animal>[] = animals.map(data => ({ data }))

const dogHolders: Holder<Dog> = groupBy(animalHolders, it => it.data.type) // Any way of doing this?

Playground link

Thanks for the help.

CodePudding user response：

Nice question...

Let's first create some utility types:

type KeysOfType<O, T> = {
  [K in keyof O]: O[K] extends T ? K : never;
}[keyof O];

This extracts all the keys of O that point to a value of type T. This is going to be used to constrain the types of discriminant to those that are of string type. They're going to be used as keys in your output type so we're not really interested in allowing discriminants of other types.

Let's also add Expand<T> to make our resulting types look nicer in intellisense.

type Expand<T> = T extends infer O ? { [K in keyof O]: O[K] } : never;

Now, let's create a type that represents the return-type of your groupBy function:

type Return<T, K extends KeysOfType<T, string>> = 
    { [KK in string & T[K]]: { [_ in K]: KK } & T }

or, optionally, liberally apply the Expand<T> type above for nicer intellisense for the consumer:

type Return<T, K extends KeysOfType<T, string>> = 
    Expand<{ [KK in string & T[K]]: Expand<{ [_ in K]: KK } & T> }>    

so now we can declare the function:

function groupBy<T, K extends KeysOfType<T, string>>(
    arr: T[], 
    keyExtractor: (element: T) => T[K]): Return<T, K>{
    throw Error();
}

and call it:

const groups = groupBy(animals, e => e.type)

for full type-safety, no matter which discriminator property is selected.

Playground Link

CodePudding user response：

There's a fairly simple solution that uses the fact that conditional types distribute over unions to get rid of the non-matching alternatives:

type GroupBy<T extends Record<D, PropertyKey>, D extends keyof T> =
  {[K in T[D]]: T extends Record<D, K> ?  T[] : never}

declare function groupBy<T extends Record<D, PropertyKey>, D extends keyof T>
  (arr: T[], keyExtractor: (element: T) => T[D]): GroupBy<T, D>

It works for the original example, but also for other discriminating keys:

interface Orange { color: 'orange', juiceContent: number }
interface Banana { color: 'yellow', length: number}
type Fruit = Orange | Banana
const fruits: Fruit[] = [{color: 'orange', juiceContent: 250},{ color: 'yellow', length: 20}]
const fruitsByColor = groupBy(fruits, it => it.color)
const yellows = fruitsByColor.yellow
// const yellows: Banana[]

TypeScript playground