Add shouldCache option

index.ts

@@ -14,6 +14,18 @@ export type CacheStorage<KeyType, ValueType> = {
 	clear?: () => unknown;
 };
 
+/**
+Controls whether a fulfilled value should be written to the cache.
+*/
+export type ShouldCache<
+	KeyType,
+	ValueType,
+	ArgumentsList extends readonly unknown[] = readonly unknown[],
+> = (
+	value: ValueType,
+	context: {key: KeyType; argumentsList: ArgumentsList},
+) => boolean | Promise<boolean>;
+
 export type Options<
 	FunctionToMemoize extends AnyAsyncFunction,
 	CacheKeyType,
@@ -52,6 +64,38 @@ export type Options<
 	@example new WeakMap()
 	*/
 	readonly cache?: CacheStorage<CacheKeyType, AsyncReturnType<FunctionToMemoize>> | false;
+
+	/**
+	Controls whether a fulfilled value should be written to the cache.
+
+	It runs after the function fulfills and before `cache.set`.
+
+	- Omit to keep current behavior (always write).
+	- Return `false` to skip writing to the cache (in-flight de-duplication is still cleared).
+	- Throw or reject to propagate the error and skip caching.
+
+	@example
+	```
+	import pMemoize from 'p-memoize';
+
+	// Only cache defined values
+	const getMaybe = pMemoize(async key => db.get(key), {
+		shouldCache: value => value !== undefined,
+	});
+
+	// Only cache non-empty arrays
+	const search = pMemoize(async query => fetchResults(query), {
+		shouldCache: value => Array.isArray(value) && value.length > 0,
+	});
+	```
+
+	Note: This only affects writes; reads from the cache are unchanged.
+	*/
+	readonly shouldCache?: ShouldCache<
+		CacheKeyType,
+		AsyncReturnType<FunctionToMemoize>,
+		Parameters<FunctionToMemoize>
+	>;
 };
 
 /**
@@ -118,7 +162,13 @@ export default function pMemoize<
 					return result;
 				} finally {
 					if (cache) {
-						await cache.set(key, result);
+						const allow = options?.shouldCache
+							? await options.shouldCache(result, {key, argumentsList: arguments_})
+							: true;
+
+						if (allow) {
+							await cache.set(key, result);
+						}
 					}
 				}
 			} finally {

readme.md

@@ -75,6 +75,34 @@ Use a different cache storage. Must implement the following methods: `.has(key)`
 
 See the [caching strategy](https://github.com/sindresorhus/mem#caching-strategy) section in the `mem` package for more information.
 
+##### shouldCache
+
+Type: `(value, {key, argumentsList}) => boolean | Promise<boolean>`
+
+Controls whether a fulfilled value should be written to the cache.
+
+It runs after the function fulfills and before `cache.set`.
+
+- Omit to keep current behavior (always write).
+- Return `false` to skip writing to the cache (in-flight de-duplication is still cleared).
+- Throw or reject to propagate the error and skip caching.
+
+```js
+import pMemoize from 'p-memoize';
+
+// Only cache defined values
+const getMaybe = pMemoize(async key => db.get(key), {
+	shouldCache: value => value !== undefined,
+});
+
+// Only cache non-empty arrays
+const search = pMemoize(async query => fetchResults(query), {
+	shouldCache: value => Array.isArray(value) && value.length > 0,
+});
+```
+
+Note: Affects only writes; reads from the cache are unchanged.
+
 ### pMemoizeDecorator(options)
 
 Returns a [decorator](https://github.com/tc39/proposal-decorators) to memoize class methods or static class methods.

test-d/index.test-d.ts

@@ -40,6 +40,32 @@ pMemoize(async () => 1, {
 	},
 });
 
+// `shouldCache` tests.
+pMemoize(async (text: string) => Boolean(text), {
+	shouldCache(value, context) {
+		expectType<boolean>(value);
+		expectType<[string]>(context.argumentsList);
+		expectType<string>(context.key);
+		return true;
+	},
+});
+
+pMemoize(async (_input: {key: string}) => 5 as const, {
+	cacheKey(_arguments: [{key: string}]): Date {
+		return new Date();
+	},
+	async shouldCache(value, context) {
+		expectType<5>(value);
+		expectType<[
+			{
+				key: string;
+			},
+		]>(context.argumentsList);
+		expectType<Date>(context.key);
+		return true;
+	},
+});
+
 // Ensures that the various cache functions infer their arguments type from the return type of `cacheKey`
 pMemoize(async (_arguments: {key: string}) => 1, {
 	cacheKey(arguments_: [{key: string}]) {

test.ts

@@ -297,6 +297,246 @@ test('.pMemoizeClear() throws when called on a disabled cache', t => {
 	});
 });
 
+test.serial('shouldCache: skips cache.set when predicate false', async t => {
+	let calls = 0;
+	const cache = new Map<unknown, unknown>();
+	const setSpy: Array<{key: unknown; value: unknown}> = [];
+	const wrappedCache = {
+		async has(key: unknown) {
+			return cache.has(key);
+		},
+		async get(key: unknown) {
+			return cache.get(key);
+		},
+		async set(key: unknown, value: unknown) {
+			setSpy.push({key, value});
+			cache.set(key, value);
+		},
+		delete(key: unknown) {
+			cache.delete(key);
+		},
+		clear() {
+			cache.clear();
+		},
+	} as const;
+
+	const memoized = pMemoize(async () => ++calls, {
+		cache: wrappedCache,
+		shouldCache() {
+			return false;
+		},
+	});
+
+	// @ts-expect-error Argument type does not match
+	t.is(await memoized('a'), 1);
+	t.is(setSpy.length, 0, 'cache.set was not called');
+	// Next call should recompute since we skipped caching
+	// @ts-expect-error Argument type does not match
+	t.is(await memoized('a'), 2);
+});
+
+test.serial('shouldCache: still de-dupes in-flight', async t => {
+	const {promise, resolve} = pDefer<number>();
+	let invocationsCount = 0;
+	const memoized = pMemoize(async () => {
+		invocationsCount++;
+		return promise;
+	}, {
+		shouldCache() {
+			return false;
+		},
+	});
+
+	const promise1 = memoized();
+	const promise2 = memoized();
+	t.is(promise1, promise2, 'in-flight promise is shared');
+	resolve(42);
+	t.is(await promise1, 42);
+	t.is(await promise2, 42);
+	t.is(invocationsCount, 1);
+});
+
+test.serial('shouldCache: not called when cache disabled', async t => {
+	let called = 0;
+	let i = 0;
+	const memoized = pMemoize(async () => ++i, {
+		cache: false,
+		shouldCache() {
+			called++;
+			return true;
+		},
+	});
+
+	// @ts-expect-error Argument type does not match
+	t.is(await memoized('x'), 1);
+	// @ts-expect-error Argument type does not match
+	t.is(await memoized('x'), 2);
+	t.is(called, 0);
+});
+
+test.serial('shouldCache: not called on rejection and nothing cached', async t => {
+	const setSpy: unknown[] = [];
+	const cache = new Map<unknown, unknown>();
+	const wrappedCache = {
+		async has(key: unknown) {
+			return cache.has(key);
+		},
+		async get(key: unknown) {
+			return cache.get(key);
+		},
+		async set(key: unknown, value: unknown) {
+			setSpy.push([key, value]);
+			cache.set(key, value);
+		},
+	} as const;
+
+	let called = 0;
+	const memoized = pMemoize(async () => {
+		throw new Error('boom');
+	}, {
+		cache: wrappedCache,
+		shouldCache() {
+			called++;
+			return true;
+		},
+	});
+
+	// @ts-expect-error Argument type does not match
+	await t.throwsAsync(async () => memoized('x'), {message: 'boom'});
+	t.is(called, 0);
+	t.is(setSpy.length, 0);
+});
+
+test.serial('shouldCache: supports async predicate', async t => {
+	const setSpy: Array<{key: unknown; value: unknown}> = [];
+	const cache = new Map<unknown, unknown>();
+	const wrappedCache = {
+		async has(key: unknown) {
+			return cache.has(key);
+		},
+		async get(key: unknown) {
+			return cache.get(key);
+		},
+		async set(key: unknown, value: unknown) {
+			setSpy.push({key, value});
+			cache.set(key, value);
+		},
+		delete(key: unknown) {
+			cache.delete(key);
+		},
+	} as const;
+
+	const memoized = pMemoize(async () => 'ok', {
+		cache: wrappedCache,
+		async shouldCache(value) {
+			await Promise.resolve();
+			return value === 'ok';
+		},
+	});
+
+	// @ts-expect-error Argument type does not match
+	t.is(await memoized('x'), 'ok');
+	t.is(setSpy.length, 1, 'cache.set was called once');
+});
+
+test.serial('shouldCache: works with custom cacheKey', async t => {
+	const setKeys: unknown[] = [];
+	const cache = new Map<unknown, unknown>();
+	const wrappedCache = {
+		async has(key: unknown) {
+			return cache.has(key);
+		},
+		async get(key: unknown) {
+			return cache.get(key);
+		},
+		async set(key: unknown, value: unknown) {
+			setKeys.push(key);
+			cache.set(key, value);
+		},
+		delete(key: unknown) {
+			cache.delete(key);
+		},
+	} as const;
+
+	const memoized = pMemoize(async () => 'v', {
+		cache: wrappedCache,
+		cacheKey: ([first]) => String(first),
+		shouldCache(_value, {key, argumentsList}) {
+			t.is(key, String(argumentsList[0]));
+			return true;
+		},
+	});
+
+	// @ts-expect-error Argument type does not match
+	await memoized(1);
+	t.deepEqual(setKeys, ['1']);
+});
+
+test.serial('shouldCache: predicate error bubbles and does not cache', async t => {
+	const cache = new Map<unknown, unknown>();
+	const setSpy: unknown[] = [];
+	const wrappedCache = {
+		async has(key: unknown) {
+			return cache.has(key);
+		},
+		async get(key: unknown) {
+			return cache.get(key);
+		},
+		async set(key: unknown, value: unknown) {
+			setSpy.push([key, value]);
+			cache.set(key, value);
+		},
+		delete(key: unknown) {
+			cache.delete(key);
+		},
+	} as const;
+
+	const memoized = pMemoize(async () => 'x', {
+		cache: wrappedCache,
+		shouldCache() {
+			throw new Error('nope');
+		},
+	});
+
+	// @ts-expect-error Argument type does not match
+	await t.throwsAsync(async () => memoized('k'), {message: 'nope'});
+	t.is(setSpy.length, 0);
+});
+
+test.serial('shouldCache: interop with external eviction', async t => {
+	// Simulate external TTL eviction by deleting after set.
+	const cache = new Map<unknown, unknown>();
+	const wrappedCache = {
+		async has(key: unknown) {
+			return cache.has(key);
+		},
+		async get(key: unknown) {
+			return cache.get(key);
+		},
+		async set(key: unknown, value: unknown) {
+			cache.set(key, value);
+		},
+		delete(key: unknown) {
+			cache.delete(key);
+		},
+	} as const;
+
+	let index = 0;
+	const memoized = pMemoize(async () => ++index, {
+		cache: wrappedCache,
+		shouldCache() {
+			return true;
+		},
+	});
+
+	// @ts-expect-error Argument type does not match
+	t.is(await memoized('a'), 1);
+	// Evict externally
+	wrappedCache.delete('a');
+	// @ts-expect-error Argument type does not match
+	t.is(await memoized('a'), 2);
+});
+
 {
 	pMemoize(async (url: string): Promise<string> => '', {
 		cacheKey: ([url]) => url,