2

The function will return an object of type {[key: string] : T[]} But I am looking to improve the type safety of this function so that in the IDE I get IntelliSense on what keys are valid in the resulting object.

I've been trying to get the keys as a const value and use that to type the object, but I can't figure out the correct syntax to get it working, nor do I know if that is even a correct approach.

export function groupBy<T>(array: T[], getKey: (item: T) => string) {
    const map = new Map<string, T[]>();
    array.forEach((item) => {
        const key = getKey(item);
        const group = map.get(key);
        group ? group.push(item) : map.set(key, [item]);
    });

    return Object.fromEntries(map.entries());
}
isherwood
  • 58,414
  • 16
  • 114
  • 157
Jacob Gad
  • 21
  • 2

1 Answers1

1

For this to work you want groupBy() to be generic not only in T, the element type of array, but also in K, the return type of getKey(). Like this:

function groupBy<T, K extends string>(array: T[], getKey: (item: T) => K) {
  const map = new Map<string, T[]>();
  array.forEach((item) => {
    const key = getKey(item);
    const group = map.get(key);
    group ? group.push(item) : map.set(key, [item]);
  });

  return Object.fromEntries(map.entries()) as { [P in K]: T[] }
}

Here K is constrained to string, so the compiler will complain if getKey() returns a non-string. But now when you call groupBy(), the compiler will try to infer K to be a union of string literal types corresponding to the particular values that getKey() can return.

Then the return type of groupBy() is { [P in K]: T[] } (equivalent to Record<K, T[]> using the Record utility type instead of { [k: string]: T[] }. That means it has keys of type K and values of type T[]. Note that I needed to assert that the value returned by Object.entries() is of that type, since the TypeScript compiler can't verify that (see Preserve Type when using Object.entries or other questions for more information).


Let's test it out with an example:

interface Foo {
  name: string;
  age: number;
}

const foos: Foo[] = [
  { name: "Alice", age: 30 },
  { name: "Bob", age: 35 },
  { name: "Carol", age: 40 },
  { name: "Dave", age: 50 }
];

const groups = groupBy(
  foos,
  (foo) => foo.age >= 40 ? "older" : foo.name.length <= 4 ? "shortname" : "other"
);
/* const groups: {
    older: Foo[];
    shortname: Foo[];
    other: Foo[];
} */

console.log("Older: " + groups.older.map(foo => foo.name).join(", "));
// Older: Carol, Dave

console.log("Shortname: " + groups.shortname.map(foo => foo.name).join(", "))
// Shortname: Bob

groups.randomKey; // error

Looks good. The compiler infers that K is the union "older" | "shortname" | "other" and therefore that the returned groups value is of type {older: Foo[]; shortname: Foo[]; other: Foo[]}. And so it lets you index into groups.older and groups.shortname but it wouldn't let you index into groups.randomKey, which is I think what your desired behavior is.

Playground link to code

jcalz
  • 264,269
  • 27
  • 359
  • 360