diff --git a/__tests__/collection.test.ts b/__tests__/collection.test.ts index 353980e..0de9655 100644 --- a/__tests__/collection.test.ts +++ b/__tests__/collection.test.ts @@ -435,3 +435,28 @@ describe('random thisArg tests', () => { }, array); }); }); + +describe('ensure() tests', () => { + function createTestCollection() { + return new Collection([ + ['a', 1], + ['b', 2], + ]); + } + + test('set new value if key does not exist', () => { + const coll = createTestCollection(); + coll.ensure('c', () => 3); + expect(coll.size).toStrictEqual(3); + expect(coll.get('c')).toStrictEqual(3); + }); + + test('return existing value if key exists', () => { + const coll = createTestCollection(); + const ensureB = coll.ensure('b', () => 3); + const getB = coll.get('b'); + expect(ensureB).toStrictEqual(2); + expect(getB).toStrictEqual(2); + expect(coll.size).toStrictEqual(2); + }); +}); diff --git a/src/index.ts b/src/index.ts index 7298809..20d9fdf 100644 --- a/src/index.ts +++ b/src/index.ts @@ -25,6 +25,22 @@ export interface Collection extends Map { export class Collection extends Map { public static readonly default: typeof Collection = Collection; + /** + * Obtains the value of the given key if it exists, otherwise sets and returns the value provided by the default value generator. + * + * @param key The key to get if it exists, or set otherwise + * @param defaultValueGenerator A function that generates the default value + * + * @example + * collection.ensure(guildId, () => defaultGuildConfig); + */ + public ensure(key: K, defaultValueGenerator: (key: K, collection: this) => V): V { + if (this.has(key)) return this.get(key)!; + const defaultValue = defaultValueGenerator(key, this); + this.set(key, defaultValue); + return defaultValue; + } + /** * Checks if all of the elements exist in the collection. *