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