Feature/improve group by key types

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"