|
| 1 | +import { logger } from '../logger'; |
| 2 | +import { isDevEnv } from './values'; |
| 3 | + |
| 4 | +/** |
| 5 | + * Iterate over an array, yielding multiple items at a time. If the size of the given array |
| 6 | + * is not divisible by the given batch size, then the length of the last items returned will |
| 7 | + * be smaller than the given batch size, i.e.: |
| 8 | + * ```typescript |
| 9 | + * items.length % batchSize |
| 10 | + * ``` |
| 11 | + * @param items - The array to iterate over. |
| 12 | + * @param batchSize - Maximum number of items to return at a time. |
| 13 | + * @param printBenchmark - If we should print benchmark of items per second |
| 14 | + */ |
| 15 | +export function* batchIterate<T>( |
| 16 | + items: T[], |
| 17 | + batchSize: number, |
| 18 | + printBenchmark = isDevEnv |
| 19 | +): Generator<T[]> { |
| 20 | + if (items.length === 0) return; |
| 21 | + const startTime = Date.now(); |
| 22 | + for (let i = 0; i < items.length; ) { |
| 23 | + const itemsRemaining = items.length - i; |
| 24 | + const sliceSize = Math.min(batchSize, itemsRemaining); |
| 25 | + yield items.slice(i, i + sliceSize); |
| 26 | + i += sliceSize; |
| 27 | + } |
| 28 | + if (printBenchmark) { |
| 29 | + const itemsPerSecond = Math.round((items.length / (Date.now() - startTime)) * 1000); |
| 30 | + const caller = new Error().stack?.split('at ')[3].trim(); |
| 31 | + logger.debug(`Iterated ${itemsPerSecond} items/second at ${caller}`); |
| 32 | + } |
| 33 | +} |
| 34 | + |
| 35 | +/** |
| 36 | + * Iterate over an `AsyncIterable`, yielding multiple items at a time. If the size of the given |
| 37 | + * array is not divisible by the given batch size, then the length of the last items returned will |
| 38 | + * be smaller than the given batch size. |
| 39 | + * |
| 40 | + * @param items - AsyncIterable |
| 41 | + * @param batchSize - Batch size |
| 42 | + * @param printBenchmark - If we should print benchmark of items per second |
| 43 | + */ |
| 44 | +export async function* asyncBatchIterate<T>( |
| 45 | + items: AsyncIterable<T>, |
| 46 | + batchSize: number, |
| 47 | + printBenchmark = isDevEnv |
| 48 | +): AsyncGenerator<T[], void, unknown> { |
| 49 | + const startTime = Date.now(); |
| 50 | + let itemCount = 0; |
| 51 | + let itemBatch: T[] = []; |
| 52 | + for await (const item of items) { |
| 53 | + itemBatch.push(item); |
| 54 | + itemCount++; |
| 55 | + if (itemBatch.length >= batchSize) { |
| 56 | + yield itemBatch; |
| 57 | + itemBatch = []; |
| 58 | + if (printBenchmark) { |
| 59 | + const itemsPerSecond = Math.round((itemCount / (Date.now() - startTime)) * 1000); |
| 60 | + const caller = new Error().stack?.split('at ')[3].trim(); |
| 61 | + logger.debug(`Iterated ${itemsPerSecond} items/second at ${caller}`); |
| 62 | + } |
| 63 | + } |
| 64 | + } |
| 65 | + if (itemBatch.length > 0) { |
| 66 | + yield itemBatch; |
| 67 | + } |
| 68 | +} |
| 69 | + |
| 70 | +/** |
| 71 | + * Convert an `AsyncIterable` to a generator |
| 72 | + * @param iter - AsyncIterable |
| 73 | + */ |
| 74 | +export async function* asyncIterableToGenerator<T>(iter: AsyncIterable<T>) { |
| 75 | + for await (const entry of iter) { |
| 76 | + yield entry; |
| 77 | + } |
| 78 | +} |
0 commit comments