Repository: marpple/FxTS (ts)

Issue:
Improve the type inference of the `groupBy` and `indexBy` functions
## ⭐ Suggestion
It looks like it could be made better by improving the type inference of the `groupBy` function.

## 💻 Use Cases

```ts
const list = [
  { type: "a", value: 1 },
  { type: "b", value: 2 },
  { type: "c", value: 3 },
  { type: "a", value: 4 } ,
  { type: "b", value: 5 },
  { type: "c", value: 6 },
  { type: "a", value: 7 },
  { type: "b", value: 8 },
  { type: "c", value: 9 },
  { type: "a", value: 10 },
] as { type: "a" | "b" | "c"; value: number; }[]

const res = groupBy(a => a.type, list);

/*
AS-IS
typeof res:
{
    a: {
        type: "a" | "b" | "c";
        value: number;
    }[];
    b: {
        type: "a" | "b" | "c";
        value: number;
    }[];
    c: {
        type: "a" | "b" | "c";
        value: number;
    }[];
}
*/

/*
TO-BE
typeof res:
{
    a: {
        type: "a";
        value: number;
    }[];
    b: {
        type: "b";
        value: number;
    }[];
    c: {
        type: "c";
        value: number;
    }[];
}
*/
```

Pull request:
Fix(#197) improve the type inference of the groupBy functions
Fixes #197 


Patch:
diff --git a/src/groupBy.ts b/src/groupBy.ts
index 10a51dd4..73225b56 100644
--- a/src/groupBy.ts
+++ b/src/groupBy.ts
@@ -4,7 +4,9 @@ import reduce from "./reduce";
 import type IterableInfer from "./types/IterableInfer";
 import type iterableInfer from "./types/IterableInfer";
 import type Key from "./types/Key";
+import type { Prettify } from "./types/Prettify";
 import type ReturnValueType from "./types/ReturnValueType";
+import type { GroupBy } from "./types/groupBy";
 
 /**
  * Splits Iterable/AsyncIterable into sets, grouped by the result of running each value through `f`.
@@ -35,23 +37,22 @@ import type ReturnValueType from "./types/ReturnValueType";
  *
  * {@link https://codesandbox.io/s/fxts-groupby-v8q3b | Try It}
  */
-
 function groupBy<A, B extends Key>(
   f: (a: A) => B,
   iterable: Iterable<A>,
-): { [K in B]: A[] };
+): Prettify<GroupBy<A, B>>;
 
 function groupBy<A, B extends Key>(
   f: (a: A) => B | Promise<B>,
   iterable: AsyncIterable<A>,
-): Promise<{ [K in B]: A[] }>;
+): Promise<Prettify<GroupBy<A, B>>>;
 
 function groupBy<
   A extends Iterable<unknown> | AsyncIterable<unknown>,
   B extends Key,
 >(
   f: (a: IterableInfer<A>) => B | Promise<B>,
-): (iterable: A) => ReturnValueType<A, { [K in B]: IterableInfer<A>[] }>;
+): (iterable: A) => ReturnValueType<A, Prettify<GroupBy<IterableInfer<A>, B>>>;
 
 function groupBy<
   A extends Iterable<unknown> | AsyncIterable<unknown>,
@@ -60,16 +61,14 @@ function groupBy<
   f: (a: IterableInfer<A>) => B | Promise<B>,
   iterable?: A,
 ):
-  | { [K in B]: IterableInfer<A>[] }
-  | Promise<{ [K in B]: IterableInfer<A>[] }>
-  | ((iterable: A) => ReturnValueType<A, { [K in B]: IterableInfer<A>[] }>) {
+  | Prettify<GroupBy<A, B>>
+  | Promise<Prettify<GroupBy<A, B>>>
+  | ((iterable: A) => ReturnValueType<A, Prettify<GroupBy<A, B>>>) {
   if (iterable === undefined) {
-    return (
-      iterable: A,
-    ): ReturnValueType<A, { [K in B]: IterableInfer<A>[] }> => {
-      return groupBy(f, iterable as any) as ReturnValueType<
+    return (iterable: A): ReturnValueType<A, GroupBy<A, B>> => {
+      return groupBy(f, iterable as any) as unknown as ReturnValueType<
         A,
-        { [K in B]: IterableInfer<A>[] }
+        GroupBy<A, B>
       >;
     };
   }
@@ -86,11 +85,11 @@ function groupBy<
       },
       obj,
       iterable,
-    );
+    ) as unknown as GroupBy<A, B>;
   }
 
   if (isAsyncIterable<iterableInfer<A>>(iterable)) {
-    return reduce(
+    const reulst = reduce(
       async (group, a) => {
         const key = await f(a);
         return (group[key] || (group[key] = [])).push(a), group;
@@ -98,6 +97,7 @@ function groupBy<
       obj,
       iterable,
     );
+    return reulst as unknown as Promise<GroupBy<A, B>>;
   }
 
   throw new TypeError("'iterable' must be type of Iterable or AsyncIterable");
diff --git a/src/indexBy.ts b/src/indexBy.ts
index 7c8e1410..94ab710c 100644
--- a/src/indexBy.ts
+++ b/src/indexBy.ts
@@ -31,12 +31,24 @@ import type ReturnValueType from "./types/ReturnValueType";
 function indexBy<A, B extends Key>(
   f: (a: A) => B,
   iterable: Iterable<A>,
-): { [K in B]: A };
+): {
+  [K in B]: A extends object
+    ? {
+        [K2 in keyof A]: A[K2] extends B ? K : A[K2];
+      }
+    : A;
+};
 
 function indexBy<A, B extends Key>(
   f: (a: A) => B | Promise<B>,
   iterable: AsyncIterable<A>,
-): Promise<{ [K in B]: A }>;
+): Promise<{
+  [K in B]: A extends object
+    ? {
+        [K2 in keyof A]: A[K2] extends B ? K : A[K2];
+      }
+    : A;
+}>;
 
 function indexBy<
   A extends Iterable<unknown> | AsyncIterable<unknown>,
diff --git a/src/partition.ts b/src/partition.ts
index f385adb8..8cc5642f 100644
--- a/src/partition.ts
+++ b/src/partition.ts
@@ -78,7 +78,10 @@ function partition<A extends Iterable<unknown> | AsyncIterable<unknown>, B>(
   f: (a: IterableInfer<A>) => B,
 ): (iterable: A) => ReturnPartitionType<A>;
 
-function partition<A extends Iterable<unknown> | AsyncIterable<unknown>, B>(
+function partition<
+  A extends Iterable<unknown> | AsyncIterable<unknown>,
+  B extends boolean,
+>(
   f: (a: IterableInfer<A>) => B,
   iterable?: A,
 ):
@@ -87,10 +90,9 @@ function partition<A extends Iterable<unknown> | AsyncIterable<unknown>, B>(
   | ((iterable: A) => ReturnPartitionType<A>) {
   if (iterable === undefined) {
     return (iterable: A): ReturnPartitionType<A> => {
-      return partition(f, iterable as any) as ReturnPartitionType<A>;
+      return partition(f, iterable as any) as unknown as ReturnPartitionType<A>;
     };
   }
-
   if (isIterable<IterableInfer<A>>(iterable)) {
     const group = groupBy((a) => {
       const key = f(a);
@@ -99,12 +101,18 @@ function partition<A extends Iterable<unknown> | AsyncIterable<unknown>, B>(
       }
       return `${Boolean(key)}`;
     }, iterable);
-    return [group["true"] || [], group["false"] || []];
+    return [group["true"] || [], group["false"] || []] as unknown as [
+      IterableInfer<A>[],
+      IterableInfer<A>[],
+    ];
   }
 
   if (isAsyncIterable<IterableInfer<A>>(iterable)) {
     const group = groupBy(async (a) => `${Boolean(await f(a))}`, iterable);
-    return group.then((group) => [group["true"] || [], group["false"] || []]);
+    return group.then((group) => [
+      group["true"] || [],
+      group["false"] || [],
+    ]) as Promise<[IterableInfer<A>[], IterableInfer<A>[]]>;
   }
 
   throw new TypeError("'iterable' must be type of Iterable or AsyncIterable");
diff --git a/src/types/Prettify.ts b/src/types/Prettify.ts
new file mode 100644
index 00000000..5f29261c
--- /dev/null
+++ b/src/types/Prettify.ts
@@ -0,0 +1,3 @@
+export type Prettify<T> = NonNullable<{
+  [K in keyof T]: T[K];
+}>;
diff --git a/src/types/ReturnPartitionType.ts b/src/types/ReturnPartitionType.ts
index 38ed5824..9b9b8149 100644
--- a/src/types/ReturnPartitionType.ts
+++ b/src/types/ReturnPartitionType.ts
@@ -1,11 +1,14 @@
-import type Awaited from "./Awaited";
 import type IterableInfer from "./IterableInfer";
 import type ReturnValueType from "./ReturnValueType";
+import type { GroupBy } from "./groupBy";
 
 type ReturnPartitionType<T extends Iterable<unknown> | AsyncIterable<unknown>> =
   ReturnValueType<
     T,
-    [Awaited<IterableInfer<T>>[], Awaited<IterableInfer<T>>[]]
+    [
+      Awaited<GroupBy<IterableInfer<T>, "true" | "false">["true"]>,
+      Awaited<GroupBy<IterableInfer<T>, "true" | "false">["false"]>,
+    ]
   >;
 
 export default ReturnPartitionType;
diff --git a/src/types/groupBy.ts b/src/types/groupBy.ts
new file mode 100644
index 00000000..e3ea98d6
--- /dev/null
+++ b/src/types/groupBy.ts
@@ -0,0 +1,9 @@
+import type Key from "./Key";
+
+export type GroupBy<A, B extends Key> = {
+  [K in B]: A extends object
+    ? {
+        [K2 in keyof A]: A[K2] extends B ? K : A[K2];
+      }[]
+    : K[];
+};


Test patch:
diff --git a/type-check/groupBy.test.ts b/type-check/groupBy.test.ts
index 9fd8961a..90aa3490 100644
--- a/type-check/groupBy.test.ts
+++ b/type-check/groupBy.test.ts
@@ -30,9 +30,74 @@ const res4 = pipe(
   groupBy((a) => Promise.resolve(a.name)),
 );
 
+type Res5To8 = {
+  a: {
+    type: "a";
+    value: number;
+  }[];
+  b: {
+    type: "b";
+    value: number;
+  }[];
+  c: {
+    type: "c";
+    value: number;
+  }[];
+};
+
+type Res9To12 = { a: "a"[]; b: "b"[]; c: "c"[] };
+
+const list1 = [
+  { type: "a", value: 1 },
+  { type: "b", value: 2 },
+  { type: "c", value: 3 },
+  { type: "a", value: 4 },
+  { type: "b", value: 5 },
+  { type: "c", value: 6 },
+  { type: "a", value: 7 },
+  { type: "b", value: 8 },
+  { type: "c", value: 9 },
+  { type: "a", value: 10 },
+] as { type: "a" | "b" | "c"; value: number }[];
+const res5 = groupBy((a) => a.type, list1);
+
+const res6 = pipe(
+  list1,
+  groupBy((a) => a.type),
+); // Res
+const res7 = groupBy((a) => a.type, toAsync(list1)); // Promise<Res>
+const res8 = pipe(
+  list1,
+  toAsync,
+  groupBy((a) => a.type),
+); // Promise<Res>
+
+const list2 = ["a", "b", "c", "a", "b", "c"] as Array<"a" | "b" | "c">;
+
+const res9 = groupBy((a) => a, list2); // Res
+const res10 = pipe(
+  list2,
+  groupBy((a) => a),
+); // Res
+
+const res11 = groupBy((a) => a, toAsync(list2)); // Promise<Res>
+const res12 = pipe(
+  list2,
+  toAsync,
+  groupBy((a) => a),
+); // Promise<Res>
+
 checks([
   check<typeof res1, { [p: string]: Data[] }, Test.Pass>(),
   check<typeof res2, { [p: string]: Data[] }, Test.Pass>(),
   check<typeof res3, { [p: string]: Data[] }, Test.Pass>(),
   check<typeof res4, Promise<{ [p: string]: Data[] }>, Test.Pass>(),
+  check<typeof res5, Res5To8, Test.Pass>(),
+  check<typeof res6, Res5To8, Test.Pass>(),
+  check<typeof res7, Promise<Res5To8>, Test.Pass>(),
+  check<typeof res8, Promise<Res5To8>, Test.Pass>(),
+  check<typeof res9, Res9To12, Test.Pass>(),
+  check<typeof res10, Res9To12, Test.Pass>(),
+  check<typeof res11, Promise<Res9To12>, Test.Pass>(),
+  check<typeof res12, Promise<Res9To12>, Test.Pass>(),
 ]);
diff --git a/type-check/indexBy.test.ts b/type-check/indexBy.test.ts
index f2aa302c..157d9ca9 100644
--- a/type-check/indexBy.test.ts
+++ b/type-check/indexBy.test.ts
@@ -29,9 +29,40 @@ const res4 = pipe(
   indexBy((a) => Promise.resolve(a.id)),
 );
 
+const res5 = indexBy((a) => a.type, [
+  { type: "a", value: 1 },
+  { type: "b", value: 2 },
+  { type: "c", value: 3 },
+  { type: "a", value: 4 },
+  { type: "b", value: 5 },
+  { type: "c", value: 6 },
+  { type: "a", value: 7 },
+  { type: "b", value: 8 },
+  { type: "c", value: 9 },
+  { type: "a", value: 10 },
+] as { type: "a" | "b" | "c"; value: number }[]);
+
 checks([
   check<typeof res1, { [index: number]: Data }, Test.Pass>(),
   check<typeof res2, { [index: number]: Data }, Test.Pass>(),
   check<typeof res3, { [index: number]: Data }, Test.Pass>(),
   check<typeof res4, Promise<{ [index: number]: Data }>, Test.Pass>(),
+  check<
+    typeof res5,
+    {
+      a: {
+        type: "a";
+        value: number;
+      };
+      b: {
+        type: "b";
+        value: number;
+      };
+      c: {
+        type: "c";
+        value: number;
+      };
+    },
+    Test.Pass
+  >(),
 ]);