I am trying to create a version of Array.groupBy() that handles more cases. Specifically, I want to be able to store a single object for a given key (instead of an array of objects with Array.groupBy()) or pluck a single property from each object.
The execution of the function works well, however the types are not always inferred properly.
The type definition is the following:
declare global {
interface Array<T> {
/**
* Groups the elements of an array by a specified field name.
*
* @param fieldName - The field name to group by.
* @param valueFieldName - Optional. The field name to use as the value in the grouped result.
* If not provided, the entire object will be used as the value.
* If set to '[]', an array of objects with the same field name will be used as the value.
* @returns An object with the grouped elements.
*/
groupBy<Key extends keyof T>(
fieldName: Key,
valueFieldName?: T[Key]|'[]',
): { [key: string]: typeof valueFieldName extends '[]' ? T[] : typeof valueFieldName extends string ? T[Key] : T };
}
}
And the following calls, despite sending back correct results, are not properly typed:
// returns { John: [{ id: 1, name: 'John' }, { id: 3, name: 'John' }], Jane: [{ id: 2, name: 'Jane' }] }
// FIXME when hovering over second the inferred type is wrong:
// { [key: string]: { id: number; name: string } }
// instead of
// { [key: string]: { id: number; name: string }[] }
const second = [{ id: 1, name: 'John' }, { id: 2, name: 'Jane' }, { id: 3, name: 'John' }]
.groupBy('name', '[]');
console.log(second)
// returns { John: 28, Jane: 30 }
// FIXME when hovering over second the inferred type is wrong:
// { [key: string]: { id: number; name: string; age: number } }
// instead of
// { [key: string]: number }
const third = [{ id: 1, name: 'John', age: 25 }, { id: 2, name: 'Jane', age: 30 }, { id: 3, name: 'John', age: 28 }]
.groupBy('name', 'age');
console.log(third)
What more can I tell Typescript so that the types are inferred properly? I feel like the parts with typeof valueFieldName extends are not optimal, but they should still work in my opinion.
Or could it be that Typescript can't infer these types precisely enough? If so, what workarounds do I have?
Obligatory caveat: Readers should be aware that it is widely considered bad practice to modify or extend the prototypes of built-in objects. See Why is extending native objects a bad practice? for more information. In what follows I will describe how to give proper typings to such an extension, but this isn't an endorsement of the approach. One could also write proper typings for a standalone version of the function. I'm only using the extension example because that's how the question is phrased.
In order for the return type of groupBy() to depend on the presence or type of the argument for the valueFieldname parameter, you need to make the function generic in the type of that parameter, and fairly directly. So in addition to the generic type K of the fieldName parameter, you should make a generic type V for the valueFieldName parameter:
interface Array<T> {
groupBy<K extends keyof T, V extends undefined | '[]' | keyof T = undefined>(
fieldName: K,
valueFieldName?: V,
): { [key: string]: V extends '[]' ? T[] : V extends keyof T ? T[V] : T };
}
The K type parameter is still constrained to keyof T, and now the V type parameter is constrained to the union of acceptable value types: either the string literal type '[]', or one of the keys of T, or undefined. Note that V defaults to undefined, so if you call the method without a valueFieldName argument, V will fall back to undefined (and not to the full undefined | '[]' | keyof T constraint).
Then the return type depends on V, via a series of conditional types. If V is "[]" then you get an array of T. If V is a keyof T then you get the property type of T and the V index (that is, the indexed access type T[V]). Otherwise V is undefined and you get just T.
Let's test to make sure it works:
const first = [{ id: 1, name: 'John' }, { id: 2, name: 'Jane' }, { id: 3, name: 'John' }]
.groupBy('name');
// const first: { [key: string]: { id: number; name: string; };
const second = [{ id: 1, name: 'John' }, { id: 2, name: 'Jane' }, { id: 3, name: 'John' }]
.groupBy('name', '[]');
// const second: { [key: string]: { id: number; name: string; }[];
const third = [{ id: 1, name: 'John', age: 25 }, { id: 2, name: 'Jane', age: 30 }, { id: 3, name: 'John', age: 28 }]
.groupBy('name', 'age');
// const third: { [key: string]: number };
Looks good; the return types are what you wanted them to be.
That's basically it, although I'll explain a little more about why the approach from the example doesn't work.
First, the type in the example for valueFieldName is essentially T[K] | '[]' | undefined, and the T[K] part of that is completely inappropriate. That's saying you want valueFieldName to be the type of the property of T at the fieldName key. So in the above examples, that would be saying groupBy('age', 'name') should be rejected and groupBy('age', 25) should be accepted. That's not consistent with the implementation. If anything, it should be keyof T | '[]' | undefined.
Next, we have to have valueFieldName of generic type V extends keyof T | '[]' | undefined and not just the specific type keyof T | '[]' | undefined. With the specific type, the return type cannot possibly depend on which of those union members the caller actually selected. The example uses the typeof operator, presumably to determine this information... but typeof valueFieldName is just keyof T | '[]' | undefined, no matter what. The typeof operator does not make the function implicitly generic in its argument. So if you want typeof valueFieldName to depend on the argument passed in, then valueFieldName simply must be generic.
That should hopefully make it clear why the original example didn't work, and why the changes were made fixed it.