Add elementSelector to toMap

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

Hide file tree

Showing 2 changed files with 27 additions and 9 deletions.
diff --git a/template/implementationTemplate.ts b/template/implementationTemplate.ts
@@ -661,11 +661,17 @@ abstract class EnumerableTemplate<T> extends AsyncEnumerable<T> {
         }
     }
 
-    public async toMap<TKey>(keySelector: (element: T) => AsyncOrSync<TKey>): Promise<Map<TKey, T>> {
-        const result = new Map<TKey, T>();
+    public async toMap<TKey>(keySelector: (element: T) => AsyncOrSync<TKey>): Promise<Map<TKey, T>>;
+    public async toMap<TKey, TElement>(keySelector: (element: T) => AsyncOrSync<TKey>, elementSelector: (element: T) => AsyncOrSync<TElement>): Promise<Map<TKey, TElement>>;
+    public async toMap<TKey, TElement = T>(keySelector: (element: T) => AsyncOrSync<TKey>, elementSelector?: (element: T) => AsyncOrSync<TElement>): Promise<Map<TKey, TElement>> {
+        if (!elementSelector) {
+            elementSelector = x => x as any; // Should only be used when TElement is not given
+        }
+
+        const result = new Map<TKey, TElement>();
 
         for await (const element of this.iterable) {
-            result.set(await keySelector(element), element);
+            result.set(await keySelector(element), await elementSelector(element));
         }
 
         return result;

diff --git a/tests/toMap.spec.ts b/tests/toMap.spec.ts
@@ -5,6 +5,7 @@ interface TestCase {
     readonly input: Iterable<number>;
     readonly output: Map<string, number>;
     readonly keySelector: (x: number) => string;
+    readonly elementSelector?: (x: number) => number;
 }
 
 describe("LINQ", () => {
@@ -15,23 +16,34 @@ describe("LINQ", () => {
             output: new Map([["0", 0], ["1", 1], ["2", 2], ["3", 3]]),
             keySelector: x => x.toString(10)
         },
+        {
+            name: "should convert to map with duplicate keys with elementSelector",
+            input: [0, 1, 2, 3, 1, 2, 3, 2, 3, 3],
+            output: new Map([["0", -0], ["1", -1], ["2", -2], ["3", -3]]),
+            keySelector: x => x.toString(10),
+            elementSelector: x => -x
+        },
         {
             name: "should work with empty iterables",
             input: [],
             output: new Map(),
             keySelector: x => ""
         }
-    ])("ToDictionary", ({name, input, output, keySelector}) => {
+    ])("ToDictionary", ({name, input, output, keySelector, elementSelector}) => {
         it(`${name} synchronously`, () => {
-            expect(from(input).toMap(keySelector)).toEqual(output);
+            expect(from(input).toMap(keySelector, elementSelector!)).toEqual(output);
         });
 
-        it(`${name} asynchronously with synchronous key selector`, async () => {
-            expect(await from(input).asAsync().toMap(keySelector)).toEqual(output);
+        it(`${name} asynchronously with synchronous selectors`, async () => {
+            expect(await from(input).asAsync().toMap(keySelector, elementSelector!)).toEqual(output);
         });
 
-        it(`${name} asynchronously with asynchronous key selector`, async () => {
-            expect(await from(input).asAsync().toMap(x => Promise.resolve(keySelector(x)))).toEqual(output);
+        it(`${name} asynchronously with asynchronous selectors`, async () => {
+            if (elementSelector) {
+                expect(await from(input).asAsync().toMap(x => Promise.resolve(keySelector(x)), x => Promise.resolve(elementSelector(x)))).toEqual(output);
+            } else {
+                expect(await from(input).asAsync().toMap(x => Promise.resolve(keySelector(x)))).toEqual(output);
+            }
         });
     });
 });