Merge pull request #2 from jsr-core/allow-malformed-iterator-on-chain

feat: support chaining malformed iterables

Author: lambdalisue

async/chain.ts

@@ -12,13 +12,33 @@
  * const iter = chain([1, 2], [3, 4]);
  * console.log(await toArray(iter)); // [1, 2, 3, 4]
  * ```
+ *
+ * It supports chaining malformed iterables.
+ *
+ * @example
+ * ```ts
+ * import { toArray } from "@core/iterutil/async/to-array";
+ * import { chain } from "@core/iterutil/async/chain";
+ *
+ * const iter = chain([1, 2], ["a", "b"], [true]);
+ * console.log(await toArray(iter)); // [1, 2, "a", "b", true]
+ * ```
  */
-export async function* chain<T>(
-  ...iterables: (Iterable<T> | AsyncIterable<T>)[]
-): AsyncIterable<T> {
+export async function* chain<
+  T extends (Iterable<unknown> | AsyncIterable<unknown>)[],
+>(
+  ...iterables: T
+): AsyncIterable<Chain<T>> {
   for await (const iterable of iterables) {
     for await (const value of iterable) {
-      yield value;
+      yield value as Chain<T>;
     }
   }
 }
+
+export type Chain<T> = T extends readonly [] ? never
+  : T extends readonly [Iterable<infer U>] ? U
+  : T extends readonly [AsyncIterable<infer U>] ? U
+  : T extends readonly [Iterable<infer U>, ...infer R] ? U | Chain<R>
+  : T extends readonly [AsyncIterable<infer U>, ...infer R] ? U | Chain<R>
+  : never;

async/chain_test.ts

@@ -33,4 +33,17 @@ Deno.test("chain", async (t) => {
     assertEquals(await toArray(result), expected);
     assertType<IsExact<typeof result, AsyncIterable<number>>>(true);
   });
+
+  await t.step("with malform iterable", async () => {
+    const result = chain(
+      toAsyncIterable([1, 2]),
+      ["a", "b"],
+      toAsyncIterable([true]),
+    );
+    const expected = [1, 2, "a", "b", true];
+    assertEquals(await toArray(result), expected);
+    assertType<
+      IsExact<typeof result, AsyncIterable<number | string | boolean>>
+    >(true);
+  });
 });

chain.ts

@@ -11,9 +11,26 @@
  * const iter = chain([1, 2], [3, 4]);
  * console.log([...iter]); // [1, 2, 3, 4]
  * ```
+ *
+ * It supports chaining malformed iterables.
+ *
+ * @example
+ * ```ts
+ * import { chain } from "@core/iterutil/chain";
+ *
+ * const iter = chain([1, 2], ["a", "b"], [true]);
+ * console.log([...iter]); // [1, 2, "a", "b", true]
+ * ```
  */
-export function* chain<T>(...iterables: Iterable<T>[]): Iterable<T> {
+export function* chain<T extends Iterable<unknown>[]>(
+  ...iterables: T
+): Iterable<Chain<T>> {
   for (const iterable of iterables) {
-    yield* iterable;
+    yield* iterable as Iterable<Chain<T>>;
   }
 }
+
+export type Chain<T> = T extends readonly [] ? never
+  : T extends readonly [Iterable<infer U>] ? U
+  : T extends readonly [Iterable<infer U>, ...infer R] ? U | Chain<R>
+  : never;

chain_test.ts

@@ -2,9 +2,20 @@ import { assertEquals } from "@std/assert";
 import { assertType, type IsExact } from "@std/testing/types";
 import { chain } from "./chain.ts";
 
-Deno.test("chain", () => {
-  const result = chain([1, 2], [3, 4], [5]);
-  const expected = [1, 2, 3, 4, 5];
-  assertEquals([...result], expected);
-  assertType<IsExact<typeof result, Iterable<number>>>(true);
+Deno.test("chain", async (t) => {
+  await t.step("uniform iterables", () => {
+    const result = chain([1, 2], [3, 4], [5]);
+    const expected = [1, 2, 3, 4, 5];
+    assertEquals([...result], expected);
+    assertType<IsExact<typeof result, Iterable<number>>>(true);
+  });
+
+  await t.step("malform iterables", () => {
+    const result = chain([1, 2], ["a", "b"], [true]);
+    const expected = [1, 2, "a", "b", true];
+    assertEquals([...result], expected);
+    assertType<IsExact<typeof result, Iterable<number | string | boolean>>>(
+      true,
+    );
+  });
 });