diff --git a/packages/util/__tests__/range.test.ts b/packages/util/__tests__/range.test.ts index 2b9a37e7c..1a822ed7e 100644 --- a/packages/util/__tests__/range.test.ts +++ b/packages/util/__tests__/range.test.ts @@ -2,11 +2,15 @@ import { describe, test, expect } from 'vitest'; import { range } from '../src/index.js'; describe('range', () => { - test('GIVEN valid range THEN valid array is returned', () => { - expect(range(0, 5)).toEqual([0, 1, 2, 3, 4, 5]); + test('GIVEN valid range and then valid numbers are returned', () => { + expect([...range(5)]).toEqual([0, 1, 2, 3, 4]); }); - test('GIVEN valid range with step THEN valid array is returned', () => { - expect(range(0, 10, 2)).toEqual([0, 2, 4, 6, 8, 10]); + test('GIVEN valid range with start and end THEN valid numbers are returned', () => { + expect([...range({ start: 0, end: 5 })]).toEqual([0, 1, 2, 3, 4]); + }); + + test('GIVEN valid range with start, end and step THEN valid numbers are returned', () => { + expect([...range({ start: 0, end: 11, step: 2 })]).toEqual([0, 2, 4, 6, 8, 10]); }); }); diff --git a/packages/util/src/functions/range.ts b/packages/util/src/functions/range.ts index 24c8e4b82..8262773f6 100644 --- a/packages/util/src/functions/range.ts +++ b/packages/util/src/functions/range.ts @@ -1,20 +1,61 @@ /** - * Yields the numbers in the given range as an array + * Options for creating a range + */ +export interface RangeOptions { + /** + * The end of the range (exclusive) + */ + end: number; + /** + * The start of the range (inclusive) + */ + start: number; + /** + * The amount to increment by + * + * @defaultValue `1` + */ + step?: number; +} + +/** + * A generator to yield numbers in a given range * - * @param start - The start of the range - * @param end - The end of the range (inclusive) - * @param step - The amount to increment between each number + * @remarks + * This method is end-exclusive, for example the last number yielded by `range(5)` is 4. If you + * prefer for the end to be included add 1 to the range or `end` option. + * @param range - A number representing the the range to yield (exclusive) or an object with start, end and step * @example * Basic range * ```ts - * range(3, 5); // [3, 4, 5] + * for (const number of range(5)) { + * console.log(number); + * } + * // Prints 0, 1, 2, 3, 4 * ``` * @example * Range with a step * ```ts - * range(3, 10, 2); // [3, 5, 7, 9] + * for (const number of range({ start: 3, end: 10, step: 2 })) { + * console.log(number); + * } + * // Prints 3, 5, 7, 9 * ``` */ -export function range(start: number, end: number, step = 1): number[] { - return Array.from({ length: (end - start) / step + 1 }, (_, index) => start + index * step); +export function* range(range: RangeOptions | number) { + let rangeEnd: number; + let start = 0; + let step = 1; + + if (typeof range === 'number') { + rangeEnd = range; + } else { + start = range.start; + rangeEnd = range.end; + step = range.step ?? 1; + } + + for (let index = start; index < rangeEnd; index += step) { + yield index; + } } diff --git a/packages/ws/src/ws/WebSocketManager.ts b/packages/ws/src/ws/WebSocketManager.ts index d6e1864b0..8b32cb57f 100644 --- a/packages/ws/src/ws/WebSocketManager.ts +++ b/packages/ws/src/ws/WebSocketManager.ts @@ -264,11 +264,12 @@ export class WebSocketManager extends AsyncEventEmitter { if (Array.isArray(this.options.shardIds)) { shardIds = this.options.shardIds; } else { - shardIds = range(this.options.shardIds.start, this.options.shardIds.end); + const { start, end } = this.options.shardIds; + shardIds = [...range({ start, end: end + 1 })]; } } else { const data = await this.fetchGatewayInformation(); - shardIds = range(0, (this.options.shardCount ?? data.shards) - 1); + shardIds = [...range(this.options.shardCount ?? data.shards)]; } this.shardIds = shardIds;