Add index parameter to the predicate of Where

Martin1994 · May 9, 2021 · 457227e · 457227e
1 parent 24a2c1b
commit 457227e
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 4 deletions.
diff --git a/template/implementationTemplate.ts b/template/implementationTemplate.ts
@@ -703,11 +703,13 @@ abstract class EnumerableTemplate<T> extends AsyncEnumerable<T> {
         yield* thatSet;
     }
 
-    public async *where(predicate: (element: T) => AsyncOrSync<boolean>): WrapWithAsyncEnumerable<T> {
+    public async *where(predicate: (element: T, index: number) => AsyncOrSync<boolean>): WrapWithAsyncEnumerable<T> {
+        let i = 0;
         for await (const element of this.iterable) {
-            if (await predicate(element)) {
+            if (await predicate(element, i)) {
                 yield element;
             }
+            i++;
         }
     }
 

diff --git a/tests/where.spec.ts b/tests/where.spec.ts
@@ -4,7 +4,7 @@ interface TestCase {
     readonly name: string;
     readonly input: Iterable<number>;
     readonly output: Iterable<number>;
-    readonly predicate: (x: number) => boolean;
+    readonly predicate: (x: number, i: number) => boolean;
 }
 
 describe("LINQ", () => {
@@ -14,6 +14,18 @@ describe("LINQ", () => {
             input: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
             output: [0, 2, 4, 6, 8],
             predicate: x => x % 2 === 0
+        },
+        {
+            name: "should filter the iterable with elements satisfying specified criteria using index",
+            input: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
+            output: [0, 2, 4],
+            predicate: (x, i) => x % 2 === 0 && i < 5
+        },
+        {
+            name: "should handle empty input",
+            input: [],
+            output: [],
+            predicate: x => x % 2 === 0
         }
     ])("Where", ({name, input, output, predicate}) => {
         it(`${name} synchronously`, () => {
@@ -25,7 +37,7 @@ describe("LINQ", () => {
         });
 
         it(`${name} asynchronously with asynchronous predicate`, async () => {
-            expect(await from(input).asAsync().where(x => Promise.resolve(predicate(x))).toArray()).toEqual(output);
+            expect(await from(input).asAsync().where((x, i) => Promise.resolve(predicate(x, i))).toArray()).toEqual(output);
         });
     });
 });