feat: add Collection#ensure (#52)

Co-authored-by: Antonio Román <kyradiscord@gmail.com>
Co-authored-by: muchnameless <12682826+muchnameless@users.noreply.github.com>
Co-authored-by: Rodry <38259440+ImRodry@users.noreply.github.com>
Co-authored-by: Kyran Mendoza <mendozakyran8@gmail.com>
Co-authored-by: kxmndz <71072915+kxmndz@users.noreply.github.com>

__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);
+	});
+});

src/index.ts

@@ -25,6 +25,22 @@ export interface Collection<K, V> extends Map<K, V> {
 export class Collection<K, V> extends Map<K, V> {
 	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.
 	 *