Feature/improve group by key types (remeda#143)

* improve groupBy key typings

by inferring the type of the key from the provided function; it must still extend a PropertyKey because that's what an object needs as index

* update typescript version

reverse type inference gets even better with this

* make returnType of groupBy partial

we can't really access keys directly because they might not exist at runtime, and Partial makes sure of that while still preserving the better key types. this also works well with Object.entries in TS 4.3 and above

* add an extra test for unknown properties

* use built-in PropertyKey instead of our own type

* fix types for indexed version

package.json

@@ -14,7 +14,7 @@
     "ts-jest": "^26.4.1",
     "ts-node": "^8.4.1",
     "typedoc": "^0.11.1",
-    "typescript": "^4.0.3"
+    "typescript": "^4.4.4"
   },
   "scripts": {
     "test": "jest --coverage && yarn run compile",

src/_types.ts

@@ -6,9 +6,6 @@ export type PredIndexedOptional<T, K> = (
   array?: T[]
 ) => K;
 
-/** types that may be returned by `keyof` */
-export type Key = string | number | symbol;
-
 /** Mapped type to remove optional, null, and undefined from all props */
 export type NonNull<T> = { [K in keyof T]-?: Exclude<T[K], null | undefined> };
 

src/groupBy.test.ts

@@ -1,3 +1,4 @@
+import { AssertEqual } from './_types';
 import { groupBy } from './groupBy';
 import { pipe } from './pipe';
 
@@ -43,3 +44,40 @@ describe('data last', () => {
     ).toEqual(expected);
   });
 });
+
+describe('groupBy typings', () => {
+  test('keys should be strictly inferred', () => {
+    const actual = groupBy(array, x => x.a);
+    const result: AssertEqual<keyof typeof actual, 1 | 2> = true;
+    expect(result).toEqual(true);
+  });
+  test('keys should be strictly inferred for indexed version', () => {
+    const actual = groupBy.indexed(array, x => x.a);
+    const result: AssertEqual<keyof typeof actual, 1 | 2> = true;
+    expect(result).toEqual(true);
+  });
+  test('keys should be strictly inferred for data last version', () => {
+    const actual = pipe(
+      array,
+      groupBy(x => x.a)
+    );
+    const result: AssertEqual<keyof typeof actual, 1 | 2> = true;
+    expect(result).toEqual(true);
+  });
+  test('keys of union types should not be directly accessible', () => {
+    interface MyType {
+      type: 'right' | 'wrong';
+    }
+
+    const myTypes: MyType[] = [{ type: 'right' }];
+
+    const grouped = groupBy(myTypes, item => item.type);
+
+    // @ts-expect-error we can't access keys unconditionally
+    expect(() => grouped.wrong.length).toThrow();
+    // @ts-expect-error other values are not allowed by typings
+    expect(grouped.notexists).toBeUndefined();
+    expect(grouped.wrong?.length).toBeUndefined();
+    expect(grouped.right?.length).toBe(1);
+  });
+});

src/groupBy.ts

@@ -13,14 +13,14 @@ import { NonEmptyArray, PredIndexedOptional, PredIndexed } from './_types';
  * @indexed
  * @category Array
  */
-export function groupBy<T>(
-  items: readonly T[],
-  fn: (item: T) => PropertyKey
-): Record<PropertyKey, NonEmptyArray<T>>;
+export function groupBy<Item, Key extends PropertyKey>(
+  items: readonly Item[],
+  fn: (item: Item) => Key
+): Partial<Record<Key, NonEmptyArray<Item>>>;
 
-export function groupBy<T>(
-  fn: (item: T) => PropertyKey
-): (array: readonly T[]) => Record<PropertyKey, NonEmptyArray<T>>;
+export function groupBy<Item, Key extends PropertyKey>(
+  fn: (item: Item) => Key
+): (array: readonly Item[]) => Partial<Record<Key, NonEmptyArray<Item>>>;
 
 /**
  * Splits a collection into sets, grouped by the result of running each value through `fn`.
@@ -54,13 +54,13 @@ const _groupBy = (indexed: boolean) => <T>(
 };
 
 export namespace groupBy {
-  export function indexed<T, K>(
-    array: readonly T[],
-    fn: PredIndexed<T, PropertyKey>
-  ): Record<string, NonEmptyArray<T>>;
-  export function indexed<T, K>(
-    fn: PredIndexed<T, PropertyKey>
-  ): (array: readonly T[]) => Record<string, NonEmptyArray<T>>;
+  export function indexed<Item, Key extends PropertyKey>(
+    array: readonly Item[],
+    fn: PredIndexed<Item, Key>
+  ): Partial<Record<Key, NonEmptyArray<Item>>>;
+  export function indexed<Item, Key extends PropertyKey>(
+    fn: PredIndexed<Item, Key>
+  ): (array: readonly Item[]) => Partial<Record<Key, NonEmptyArray<Item>>>;
   export function indexed() {
     return purry(_groupBy(true), arguments);
   }

src/pathOr.ts

@@ -1,5 +1,5 @@
 import { purry } from './purry';
-import { NonNull, Key } from './_types';
+import { NonNull } from './_types';
 
 /**
  * Given a union of indexable types `T`, we derive an indexable type
@@ -27,7 +27,7 @@ import { NonNull, Key } from './_types';
 type Pathable<T> = { [K in AllKeys<T>]: TypesForKey<T, K> };
 
 type AllKeys<T> = T extends infer I ? keyof I : never;
-type TypesForKey<T, K extends Key> = T extends infer I
+type TypesForKey<T, K extends PropertyKey> = T extends infer I
   ? K extends keyof I
     ? I[K]
     : never

src/reverse.test.ts

@@ -23,9 +23,9 @@ describe('data first', () => {
     });
 
     test('variadic tuples', () => {
-      const input: [number, ...Array<string>] = [1, 'two', 'three'];
+      const input: [number, ...string[]] = [1, 'two', 'three'];
       const actual = reverse(input);
-      const result: AssertEqual<typeof actual, (string | number)[]> = true;
+      const result: AssertEqual<typeof actual, [...string[], number]> = true;
       expect(result).toEqual(true);
     });
   });
@@ -54,7 +54,7 @@ describe('data last', () => {
     test('variadic tuples', () => {
       const input: [number, ...Array<string>] = [1, 'two', 'three'];
       const actual = pipe(input, reverse());
-      const result: AssertEqual<typeof actual, (string | number)[]> = true;
+      const result: AssertEqual<typeof actual, [...string[], number]> = true;
       expect(result).toEqual(true);
     });
   });

yarn.lock

@@ -3524,10 +3524,10 @@ typescript@2.7.2:
   version "2.7.2"
   resolved "https://registry.yarnpkg.com/typescript/-/typescript-2.7.2.tgz#2d615a1ef4aee4f574425cdff7026edf81919836"
 
-typescript@^4.0.3:
-  version "4.0.3"
-  resolved "https://registry.yarnpkg.com/typescript/-/typescript-4.0.3.tgz#153bbd468ef07725c1df9c77e8b453f8d36abba5"
-  integrity sha512-tEu6DGxGgRJPb/mVPIZ48e69xCn2yRmCgYmDugAVwmJ6o+0u1RI18eO7E7WBTLYLaEVVOhwQmcdhQHweux/WPg==
+typescript@^4.4.4:
+  version "4.4.4"
+  resolved "https://registry.yarnpkg.com/typescript/-/typescript-4.4.4.tgz#2cd01a1a1f160704d3101fd5a58ff0f9fcb8030c"
+  integrity sha512-DqGhF5IKoBl8WNf8C1gu8q0xZSInh9j1kJJMqT3a94w1JzVaBU4EXOSMrz9yDqMT0xt3selp83fuFMQ0uzv6qA==
 
 uglify-js@^3.1.4:
   version "3.13.5"