fix(util): update AsyncCollection with better types #2657

Author: JadenSimon

src/shared/utilities/asyncCollection.ts

@@ -9,21 +9,42 @@ import { Coalesce } from './tsUtils'
  * High-level abstraction over async generator functions of the form `async function*` {@link AsyncGenerator}
  */
 export interface AsyncCollection<T> extends AsyncIterable<T> {
-    /** Flattens the collection 1-level deep */
+    /**
+     * Flattens the collection 1-level deep.
+     */
     flatten(): AsyncCollection<SafeUnboxIterable<T>>
-    /** Applies a mapping transform to the output generator */
-    map<U>(fn: (obj: T) => U): AsyncCollection<U>
-    /** Filters out results which will _not_ be passed on to further transformations. */
-    filter<U extends T>(predicate: Predicate<T, U>): AsyncCollection<U>
-    /** Uses only the first 'count' number of values returned by the generator. */
-    take(count: number): AsyncCollection<T>
-    /** Converts the collection to a Promise, resolving to an array of all values returned by the generator. */
+
+    /**
+     * Applies a mapping transform to the output generator.
+     */
+    map<U>(fn: (obj: T) => Promise<U> | U): AsyncCollection<U>
+
+    /**
+     * Filters out results which will _not_ be passed on to further transformations.
+     */
+    filter<U extends T>(predicate: (item: T) => item is U): AsyncCollection<U>
+    filter<U extends T>(predicate: (item: T) => boolean): AsyncCollection<U>
+
+    /**
+     * Uses only the first 'count' number of values returned by the generator.
+     */
+    limit(count: number): AsyncCollection<T>
+
+    /**
+     * Converts the collection to a Promise, resolving to an array of all values returned by the generator.
+     */
     promise(): Promise<T[]>
-    /** Converts the collection to a Map, using either a property of the item or a function to select keys. */
+
+    /**
+     * Converts the collection to a Map, using either a property of the item or a function to select keys
+     */
     toMap<K extends StringProperty<T>, U extends string = never>(
         selector: KeySelector<T, U> | K
     ): Promise<Map<Coalesce<U, T[K]>, T>>
-    /** Returns an iterator directly from the underlying generator, preserving values returned. */
+
+    /**
+     * Returns an iterator directly from the underlying generator, preserving values returned.
+     */
     iterator(): AsyncIterator<T, T | void>
 }
 
@@ -55,8 +76,8 @@ export function toCollection<T>(generator: () => AsyncGenerator<T, T | undefined
         flatten: () => toCollection<SafeUnboxIterable<T>>(() => delegateGenerator(generator(), flatten)),
         filter: <U extends T>(predicate: Predicate<T, U>) =>
             toCollection<U>(() => filterGenerator<T, U>(generator(), predicate)),
-        map: <U>(fn: (item: T) => U) => toCollection<U>(() => mapGenerator(generator(), fn)),
-        take: (count: number) => toCollection(() => delegateGenerator(generator(), takeFrom(count))),
+        map: <U>(fn: (item: T) => Promise<U> | U) => toCollection<U>(() => mapGenerator(generator(), fn)),
+        limit: (count: number) => toCollection(() => delegateGenerator(generator(), takeFrom(count))),
         promise: () => promise(iterable),
         toMap: <U extends string = never, K extends StringProperty<T> = never>(selector: KeySelector<T, U> | K) =>
             asyncIterableToMap(iterable, selector),
@@ -66,7 +87,7 @@ export function toCollection<T>(generator: () => AsyncGenerator<T, T | undefined
 
 async function* mapGenerator<T, U, R = T>(
     generator: AsyncGenerator<T, R | undefined | void>,
-    fn: (item: T | R) => U
+    fn: (item: T | R) => Promise<U> | U
 ): AsyncGenerator<U, U | undefined> {
     while (true) {
         const { value, done } = await generator.next()
@@ -79,9 +100,11 @@ async function* mapGenerator<T, U, R = T>(
     }
 }
 
+type Predicate<T, U extends T> = (item: T) => item is U
+
 async function* filterGenerator<T, U extends T, R = T>(
     generator: AsyncGenerator<T, R | undefined | void>,
-    predicate: Predicate<T | R, U>
+    predicate: Predicate<T | R, U> | ((item: T | R) => boolean)
 ): AsyncGenerator<U, U | void> {
     while (true) {
         const { value, done } = await generator.next()
@@ -159,8 +182,6 @@ async function promise<T>(iterable: AsyncIterable<T>): Promise<T[]> {
     return result
 }
 
-type Predicate<T, U extends T> = ((item: T) => item is U) | ((item: T) => boolean)
-
 function addToMap<T, U extends string>(map: Map<string, T>, selector: KeySelector<T, U> | StringProperty<T>, item: T) {
     const key = typeof selector === 'function' ? selector(item) : item[selector]
     if (key) {

src/test/shared/utilities/asyncCollection.test.ts

@@ -78,16 +78,16 @@ describe('AsyncCollection', function () {
         assert.deepStrictEqual(await filtered.promise(), expected)
     })
 
-    it('can take', async function () {
-        const take = toCollection(gen).take(2)
-        assert.deepStrictEqual(await take.promise(), [0, 1])
+    it('can limit', async function () {
+        const limited = toCollection(gen).limit(2)
+        assert.deepStrictEqual(await limited.promise(), [0, 1])
     })
 
-    it('returns nothing if using non-positive count', async function () {
-        const takeZero = toCollection(gen).take(0)
-        const takeNeg1 = toCollection(gen).take(-1)
-        assert.deepStrictEqual(await takeZero.promise(), [])
-        assert.deepStrictEqual(await takeNeg1.promise(), [])
+    it('returns nothing if using non-positive limit count', async function () {
+        const limitZero = toCollection(gen).limit(0)
+        const limitNeg1 = toCollection(gen).limit(-1)
+        assert.deepStrictEqual(await limitZero.promise(), [])
+        assert.deepStrictEqual(await limitNeg1.promise(), [])
     })
 
     it('is immutable', async function () {
@@ -103,6 +103,19 @@ describe('AsyncCollection', function () {
         assert.deepStrictEqual(await z.promise(), [2, 2, 3, 3])
     })
 
+    it('can map with async functions', async function () {
+        const double = async (n: number) => 2 * n
+
+        const mapped = toCollection(gen)
+            .map(o => o + 1)
+            .map(double)
+            .map(o => o - 1)
+            .map(double)
+            .map(o => `${o}!`)
+
+        assert.deepStrictEqual(await mapped.promise(), ['2!', '6!', '10!'])
+    })
+
     it('does not iterate over the generator when applying transformations', async function () {
         let callCount = 0
         async function* count() {
@@ -117,7 +130,7 @@ describe('AsyncCollection', function () {
             .flatten()
         await new Promise(r => setImmediate(r))
         assert.strictEqual(callCount, 0)
-        assert.deepStrictEqual(await x.take(6).promise(), [2, 4, 3, 9, 4, 16])
+        assert.deepStrictEqual(await x.limit(6).promise(), [2, 4, 3, 9, 4, 16])
         assert.strictEqual(callCount, 6)
     })