feat(lib): add tuple return type inference (#37) (#38)

TomerAberbach · Nov 24, 2024 · 682d546 · 682d546
1 parent fe1f2a0
commit 682d546
Show file tree

Hide file tree

Showing 4 changed files with 581 additions and 2 deletions.
diff --git a/src/operations/filters.d.ts b/src/operations/filters.d.ts
@@ -143,6 +143,15 @@ export const filterConcur: {
  * @category Filters
  */
 export const filterMap: {
+  // These overloads help with inferring tuple types returned from the callback.
+  <From, To extends unknown[] | []>(
+    fn: (value: From) => To | null | undefined,
+  ): (iterable: Iterable<From>) => Iterable<NonNullable<To>>
+  <From, To extends unknown[] | []>(
+    fn: (value: From) => To | null | undefined,
+    iterable: Iterable<From>,
+  ): Iterable<NonNullable<To>>
+
   <From, To>(
     fn: (value: From) => To | null | undefined,
   ): (iterable: Iterable<From>) => Iterable<NonNullable<To>>
@@ -177,6 +186,15 @@ export const filterMap: {
  * @category Filters
  */
 export const filterMapAsync: {
+  // These overloads help with inferring tuple types returned from the callback.
+  <From, To extends unknown[] | []>(
+    fn: (value: From) => MaybePromiseLike<To | null | undefined>,
+  ): (asyncIterable: AsyncIterable<From>) => AsyncIterable<NonNullable<To>>
+  <From, To extends unknown[] | []>(
+    fn: (value: From) => MaybePromiseLike<To | null | undefined>,
+    asyncIterable: AsyncIterable<From>,
+  ): AsyncIterable<NonNullable<To>>
+
   <From, To>(
     fn: (value: From) => MaybePromiseLike<To | null | undefined>,
   ): (asyncIterable: AsyncIterable<From>) => AsyncIterable<NonNullable<To>>
@@ -211,6 +229,15 @@ export const filterMapAsync: {
  * @category Filters
  */
 export const filterMapConcur: {
+  // These overloads help with inferring tuple types returned from the callback.
+  <From, To extends unknown[] | []>(
+    fn: (value: From) => MaybePromiseLike<To | null | undefined>,
+  ): (concurIterable: ConcurIterable<From>) => ConcurIterable<NonNullable<To>>
+  <From, To extends unknown[] | []>(
+    fn: (value: From) => MaybePromiseLike<To | null | undefined>,
+    concurIterable: ConcurIterable<From>,
+  ): ConcurIterable<NonNullable<To>>
+
   <From, To>(
     fn: (value: From) => MaybePromiseLike<To | null | undefined>,
   ): (concurIterable: ConcurIterable<From>) => ConcurIterable<NonNullable<To>>

diff --git a/src/operations/transforms.d.ts b/src/operations/transforms.d.ts
@@ -22,6 +22,15 @@ import type { ConcurIterable } from './core.js'
  * @category Transforms
  */
 export const map: {
+  // These overloads help with inferring tuple types returned from the callback.
+  <From, To extends unknown[] | []>(
+    fn: (value: From) => To,
+  ): (iterable: Iterable<From>) => Iterable<To>
+  <From, To extends unknown[] | []>(
+    fn: (value: From) => To,
+    iterable: Iterable<From>,
+  ): Iterable<To>
+
   <From, To>(
     fn: (value: From) => To,
   ): (iterable: Iterable<From>) => Iterable<To>
@@ -49,6 +58,15 @@ export const map: {
  * @category Transforms
  */
 export const mapAsync: {
+  // These overloads help with inferring tuple types returned from the callback.
+  <From, To extends unknown[] | []>(
+    fn: (value: From) => MaybePromiseLike<To>,
+  ): (asyncIterable: AsyncIterable<From>) => AsyncIterable<To>
+  <From, To extends unknown[] | []>(
+    fn: (value: From) => MaybePromiseLike<To>,
+    asyncIterable: AsyncIterable<From>,
+  ): AsyncIterable<To>
+
   <From, To>(
     fn: (value: From) => MaybePromiseLike<To>,
   ): (asyncIterable: AsyncIterable<From>) => AsyncIterable<To>
@@ -79,6 +97,15 @@ export const mapAsync: {
  * @category Transforms
  */
 export const mapConcur: {
+  // These overloads help with inferring tuple types returned from the callback.
+  <From, To extends unknown[] | []>(
+    fn: (value: From) => MaybePromiseLike<To>,
+  ): (concurIterable: ConcurIterable<From>) => ConcurIterable<To>
+  <From, To extends unknown[] | []>(
+    fn: (value: From) => MaybePromiseLike<To>,
+    concurIterable: ConcurIterable<From>,
+  ): ConcurIterable<To>
+
   <From, To>(
     fn: (value: From) => MaybePromiseLike<To>,
   ): (concurIterable: ConcurIterable<From>) => ConcurIterable<To>
@@ -109,6 +136,15 @@ export const mapConcur: {
  * @category Transforms
  */
 export const flatMap: {
+  // These overloads help with inferring tuple types returned from the callback.
+  <From, To extends unknown[] | []>(
+    fn: (value: From) => Iterable<To>,
+  ): (iterable: Iterable<From>) => Iterable<To>
+  <From, To extends unknown[] | []>(
+    fn: (value: From) => Iterable<To>,
+    iterable: Iterable<From>,
+  ): Iterable<To>
+
   <From, To>(
     fn: (value: From) => Iterable<To>,
   ): (iterable: Iterable<From>) => Iterable<To>
@@ -140,6 +176,15 @@ export const flatMap: {
  * @category Transforms
  */
 export const flatMapAsync: {
+  // These overloads help with inferring tuple types returned from the callback.
+  <From, To extends unknown[] | []>(
+    fn: (value: From) => MaybePromiseLike<Iterable<To> | AsyncIterable<To>>,
+  ): (asyncIterable: AsyncIterable<From>) => AsyncIterable<To>
+  <From, To extends unknown[] | []>(
+    fn: (value: From) => MaybePromiseLike<Iterable<To> | AsyncIterable<To>>,
+    asyncIterable: AsyncIterable<From>,
+  ): AsyncIterable<To>
+
   <From, To>(
     fn: (value: From) => MaybePromiseLike<Iterable<To> | AsyncIterable<To>>,
   ): (asyncIterable: AsyncIterable<From>) => AsyncIterable<To>
@@ -171,6 +216,23 @@ export const flatMapAsync: {
  * @category Transforms
  */
 export const flatMapConcur: {
+  // These overloads help with inferring tuple types returned from the callback.
+  <From, To extends unknown[] | []>(
+    fn: (
+      value: From,
+    ) => MaybePromiseLike<
+      Iterable<To> | AsyncIterable<To> | ConcurIterable<To>
+    >,
+  ): (concurIterable: ConcurIterable<From>) => ConcurIterable<To>
+  <From, To extends unknown[] | []>(
+    fn: (
+      value: From,
+    ) => MaybePromiseLike<
+      Iterable<To> | AsyncIterable<To> | ConcurIterable<To>
+    >,
+    concurIterable: ConcurIterable<From>,
+  ): ConcurIterable<To>
+
   <From, To>(
     fn: (
       value: From,

diff --git a/test/operations/filters.ts b/test/operations/filters.ts
@@ -241,6 +241,32 @@ test.skip(`filterMap types are correct`, () => {
       filterMap(o => o.n),
     ),
   ).toMatchTypeOf<Iterable<number>>()
+  expectTypeOf(
+    filterMap(o => o.n, [{ n: 1 }, { n: 2 }, { m: 3 }]),
+  ).toMatchTypeOf<Iterable<number>>()
+
+  expectTypeOf(
+    pipe(
+      [{ n: 1 }, { n: 2 }, { m: 3 }],
+      filterMap(o => (o.n ? Array.from({ length: o.n }, () => 42) : null)),
+    ),
+  ).toMatchTypeOf<Iterable<number[]>>()
+  expectTypeOf(
+    filterMap(
+      o => (o.n ? Array.from({ length: o.n }, () => 42) : null),
+      [{ n: 1 }, { n: 2 }, { m: 3 }],
+    ),
+  ).toMatchTypeOf<Iterable<number[]>>()
+
+  expectTypeOf(
+    pipe(
+      [{ n: 1 }, { n: 2 }, { m: 3 }],
+      filterMap(o => (o.n ? [o.n, 1] : null)),
+    ),
+  ).toMatchTypeOf<Iterable<[number, number]>>()
+  expectTypeOf(
+    filterMap(o => (o.n ? [o.n, 1] : null), [{ n: 1 }, { n: 2 }, { m: 3 }]),
+  ).toMatchTypeOf<Iterable<[number, number]>>()
 })
 
 test.prop([
@@ -273,12 +299,70 @@ test.skip(`filterMapAsync types are correct`, () => {
       filterMapAsync(o => o.n),
     ),
   ).toMatchTypeOf<AsyncIterable<number>>()
+  expectTypeOf(
+    filterMapAsync(o => o.n, asAsync([{ n: 1 }, { n: 2 }, { m: 3 }])),
+  ).toMatchTypeOf<AsyncIterable<number>>()
   expectTypeOf(
     pipe(
       asAsync([{ n: 1 }, { n: 2 }, { m: 3 }]),
-      filterMapAsync(o => Promise.resolve(o.n)),
+      filterMapAsync(async o => o.n),
     ),
   ).toMatchTypeOf<AsyncIterable<number>>()
+  expectTypeOf(
+    filterMapAsync(async o => o.n, asAsync([{ n: 1 }, { n: 2 }, { m: 3 }])),
+  ).toMatchTypeOf<AsyncIterable<number>>()
+
+  expectTypeOf(
+    pipe(
+      asAsync([{ n: 1 }, { n: 2 }, { m: 3 }]),
+      filterMapAsync(o => (o.n ? Array.from({ length: o.n }, () => 42) : null)),
+    ),
+  ).toMatchTypeOf<AsyncIterable<number[]>>()
+  expectTypeOf(
+    filterMapAsync(
+      o => (o.n ? Array.from({ length: o.n }, () => 42) : null),
+      asAsync([{ n: 1 }, { n: 2 }, { m: 3 }]),
+    ),
+  ).toMatchTypeOf<AsyncIterable<number[]>>()
+  expectTypeOf(
+    pipe(
+      asAsync([{ n: 1 }, { n: 2 }, { m: 3 }]),
+      filterMapAsync(async o =>
+        o.n ? Array.from({ length: o.n }, () => 42) : null,
+      ),
+    ),
+  ).toMatchTypeOf<AsyncIterable<number[]>>()
+  expectTypeOf(
+    filterMapAsync(
+      async o => (o.n ? Array.from({ length: o.n }, () => 42) : null),
+      asAsync([{ n: 1 }, { n: 2 }, { m: 3 }]),
+    ),
+  ).toMatchTypeOf<AsyncIterable<number[]>>()
+
+  expectTypeOf(
+    pipe(
+      asAsync([{ n: 1 }, { n: 2 }, { m: 3 }]),
+      filterMapAsync(o => (o.n ? [o.n, 1] : null)),
+    ),
+  ).toMatchTypeOf<AsyncIterable<[number, number]>>()
+  expectTypeOf(
+    filterMapAsync(
+      o => (o.n ? [o.n, 1] : null),
+      asAsync([{ n: 1 }, { n: 2 }, { m: 3 }]),
+    ),
+  ).toMatchTypeOf<AsyncIterable<[number, number]>>()
+  expectTypeOf(
+    pipe(
+      asAsync([{ n: 1 }, { n: 2 }, { m: 3 }]),
+      filterMapAsync(async o => (o.n ? [o.n, 1] : null)),
+    ),
+  ).toMatchTypeOf<AsyncIterable<[number, number]>>()
+  expectTypeOf(
+    filterMapAsync(
+      async o => (o.n ? [o.n, 1] : null),
+      asAsync([{ n: 1 }, { n: 2 }, { m: 3 }]),
+    ),
+  ).toMatchTypeOf<AsyncIterable<[number, number]>>()
 })
 
 test.prop([
@@ -316,12 +400,78 @@ test.skip(`filterMapConcur types are correct`, () => {
       filterMapConcur(o => o.n),
     ),
   ).toMatchTypeOf<ConcurIterable<number>>()
+  expectTypeOf(
+    filterMapConcur(
+      o => o.n,
+      asConcur<{ n?: number; m?: number }>([{ n: 1 }, { n: 2 }, { m: 3 }]),
+    ),
+  ).toMatchTypeOf<ConcurIterable<number>>()
   expectTypeOf(
     pipe(
       asConcur([{ n: 1 }, { n: 2 }, { m: 3 }]),
-      filterMapConcur(o => Promise.resolve(o.n)),
+      filterMapConcur(async o => o.n),
     ),
   ).toMatchTypeOf<ConcurIterable<number>>()
+  expectTypeOf(
+    filterMapConcur(
+      async o => o.n,
+      asConcur<{ n?: number; m?: number }>([{ n: 1 }, { n: 2 }, { m: 3 }]),
+    ),
+  ).toMatchTypeOf<ConcurIterable<number>>()
+
+  expectTypeOf(
+    pipe(
+      asConcur([{ n: 1 }, { n: 2 }, { m: 3 }]),
+      filterMapConcur(o =>
+        o.n ? Array.from({ length: o.n }, () => 42) : null,
+      ),
+    ),
+  ).toMatchTypeOf<ConcurIterable<number[]>>()
+  expectTypeOf(
+    filterMapConcur(
+      o => (o.n ? Array.from({ length: o.n }, () => 42) : null),
+      asConcur<{ n?: number; m?: number }>([{ n: 1 }, { n: 2 }, { m: 3 }]),
+    ),
+  ).toMatchTypeOf<ConcurIterable<number[]>>()
+  expectTypeOf(
+    pipe(
+      asConcur([{ n: 1 }, { n: 2 }, { m: 3 }]),
+      filterMapConcur(async o =>
+        o.n ? Array.from({ length: o.n }, () => 42) : null,
+      ),
+    ),
+  ).toMatchTypeOf<ConcurIterable<number[]>>()
+  expectTypeOf(
+    filterMapConcur(
+      async o => (o.n ? Array.from({ length: o.n }, () => 42) : null),
+      asConcur<{ n?: number; m?: number }>([{ n: 1 }, { n: 2 }, { m: 3 }]),
+    ),
+  ).toMatchTypeOf<ConcurIterable<number[]>>()
+
+  expectTypeOf(
+    pipe(
+      asConcur([{ n: 1 }, { n: 2 }, { m: 3 }]),
+      filterMapConcur(o => (o.n ? [o.n, 1] : null)),
+    ),
+  ).toMatchTypeOf<ConcurIterable<[number, number]>>()
+  expectTypeOf(
+    filterMapConcur(
+      o => (o.n ? [o.n, 1] : null),
+      asConcur<{ n?: number; m?: number }>([{ n: 1 }, { n: 2 }, { m: 3 }]),
+    ),
+  ).toMatchTypeOf<ConcurIterable<[number, number]>>()
+  expectTypeOf(
+    pipe(
+      asConcur([{ n: 1 }, { n: 2 }, { m: 3 }]),
+      filterMapConcur(async o => (o.n ? [o.n, 1] : null)),
+    ),
+  ).toMatchTypeOf<ConcurIterable<[number, number]>>()
+  expectTypeOf(
+    filterMapConcur(
+      async o => (o.n ? [o.n, 1] : null),
+      asConcur<{ n?: number; m?: number }>([{ n: 1 }, { n: 2 }, { m: 3 }]),
+    ),
+  ).toMatchTypeOf<ConcurIterable<[number, number]>>()
 })
 
 test.prop([