SetOptional/SetRequired/SetReadonly: Fix instantiations with index signatures (#1014)

Author: som-sm

source/internal/object.d.ts

@@ -1,5 +1,6 @@
 import type {Simplify} from '../simplify';
 import type {UnknownArray} from '../unknown-array';
+import type {KeysOfUnion} from '../keys-of-union';
 import type {FilterDefinedKeys, FilterOptionalKeys} from './keys';
 import type {NonRecursiveType} from './type';
 import type {ToString} from './string';
@@ -80,3 +81,44 @@ export type UndefinedToOptional<T extends object> = Simplify<
 	[Key in keyof Pick<T, FilterOptionalKeys<T>>]?: Exclude<T[Key], undefined>;
 }
 >;
+
+/**
+Works similar to the built-in `Pick` utility type, except for the following differences:
+- Distributes over union types and allows picking keys from any member of the union type.
+- Primitives types are returned as-is.
+- Picks all keys if `Keys` is `any`.
+- Doesn't pick `number` from a `string` index signature.
+
+@example
+```
+type ImageUpload = {
+	url: string;
+	size: number;
+	thumbnailUrl: string;
+};
+
+type VideoUpload = {
+	url: string;
+	duration: number;
+	encodingFormat: string;
+};
+
+// Distributes over union types and allows picking keys from any member of the union type
+type MediaDisplay = HomomorphicPick<ImageUpload | VideoUpload, "url" | "size" | "duration">;
+//=> {url: string; size: number} | {url: string; duration: number}
+
+// Primitive types are returned as-is
+type Primitive = HomomorphicPick<string | number, 'toUpperCase' | 'toString'>;
+//=> string | number
+
+// Picks all keys if `Keys` is `any`
+type Any = HomomorphicPick<{a: 1; b: 2} | {c: 3}, any>;
+//=> {a: 1; b: 2} | {c: 3}
+
+// Doesn't pick `number` from a `string` index signature
+type IndexSignature = HomomorphicPick<{[k: string]: unknown}, number>;
+//=> {}
+*/
+export type HomomorphicPick<T, Keys extends KeysOfUnion<T>> = {
+	[P in keyof T as Extract<P, Keys>]: T[P]
+};

source/set-optional.d.ts

@@ -1,4 +1,6 @@
 import type {Except} from './except';
+import type {HomomorphicPick} from './internal';
+import type {KeysOfUnion} from './keys-of-union';
 import type {Simplify} from './simplify';
 
 /**
@@ -32,6 +34,6 @@ export type SetOptional<BaseType, Keys extends keyof BaseType> =
 		// Pick just the keys that are readonly from the base type.
 		Except<BaseType, Keys> &
 		// Pick the keys that should be mutable from the base type and make them mutable.
-		Partial<Except<BaseType, Exclude<keyof BaseType, Keys>>>
+		Partial<HomomorphicPick<BaseType, Keys & KeysOfUnion<BaseType>>>
 		>
 		: never;

source/set-readonly.d.ts

@@ -1,4 +1,6 @@
 import type {Except} from './except';
+import type {HomomorphicPick} from './internal';
+import type {KeysOfUnion} from './keys-of-union';
 import type {Simplify} from './simplify';
 
 /**
@@ -33,6 +35,6 @@ export type SetReadonly<BaseType, Keys extends keyof BaseType> =
 	BaseType extends unknown
 		? Simplify<
 		Except<BaseType, Keys> &
-		Readonly<Except<BaseType, Exclude<keyof BaseType, Keys>>>
+		Readonly<HomomorphicPick<BaseType, Keys & KeysOfUnion<BaseType>>>
 		>
 		: never;

source/set-required.d.ts

@@ -1,4 +1,6 @@
 import type {Except} from './except';
+import type {HomomorphicPick} from './internal';
+import type {KeysOfUnion} from './keys-of-union';
 import type {Simplify} from './simplify';
 
 /**
@@ -35,6 +37,6 @@ export type SetRequired<BaseType, Keys extends keyof BaseType> =
 		// Pick just the keys that are optional from the base type.
 		Except<BaseType, Keys> &
 		// Pick the keys that should be required from the base type and make them required.
-		Required<Except<BaseType, Exclude<keyof BaseType, Keys>>>
+		Required<HomomorphicPick<BaseType, Keys & KeysOfUnion<BaseType>>>
 		>
 		: never;

test-d/distributed-pick.ts

@@ -76,3 +76,14 @@ if (pickedUnion.discriminant === 'A') {
 	// @ts-expect-error
 	const _bar = pickedUnion.bar; // eslint-disable-line @typescript-eslint/no-unsafe-assignment
 }
+
+// Preserves property modifiers
+declare const test1: DistributedPick<{readonly 'a': 1; 'b'?: 2; readonly 'c'?: 3}, 'a' | 'b' | 'c'>;
+expectType<{readonly 'a': 1; 'b'?: 2; readonly 'c'?: 3}>(test1);
+
+declare const test2: DistributedPick<{readonly 'a': 1; 'b'?: 2} | {readonly 'c'?: 3}, 'a' | 'b' | 'c'>;
+expectType<{readonly 'a': 1; 'b'?: 2} | {readonly 'c'?: 3}>(test2);
+
+// Works with index signatures
+declare const test4: DistributedPick<{[k: string]: unknown; a?: number; b: string}, 'a' | 'b'>;
+expectType<{a?: number; b: string}>(test4);

test-d/internal/homomorphic-pick.ts

@@ -0,0 +1,53 @@
+import {expectType} from 'tsd';
+import type {HomomorphicPick} from '../../source/internal';
+
+// Picks specified keys
+declare const test1: HomomorphicPick<{a: 1; b: 2; c: 3}, 'a' | 'b'>;
+expectType<{a: 1; b: 2}>(test1);
+
+// Works with unions
+declare const test2: HomomorphicPick<{a: 1; b: 2} | {a: 3; c: 4}, 'a'>;
+expectType<{a: 1} | {a: 3}>(test2);
+
+declare const test3: HomomorphicPick<{a: 1; b: 2} | {c: 3; d: 4}, 'a' | 'c'>;
+expectType<{a: 1} | {c: 3}>(test3);
+
+// Preserves property modifiers
+declare const test4: HomomorphicPick<{readonly a: 1; b?: 2; readonly c?: 3}, 'a' | 'c'>;
+expectType<{readonly a: 1; readonly c?: 3}>(test4);
+
+declare const test5: HomomorphicPick<{readonly a: 1; b?: 2} | {readonly c?: 3; d?: 4}, 'a' | 'c'>;
+expectType<{readonly a: 1} | {readonly c?: 3}>(test5);
+
+// Passes through primitives unchanged
+declare const test6: HomomorphicPick<string, never>;
+expectType<string>(test6);
+
+declare const test7: HomomorphicPick<number, never>;
+expectType<number>(test7);
+
+declare const test8: HomomorphicPick<boolean, never>;
+expectType<boolean>(test8);
+
+declare const test9: HomomorphicPick<bigint, never>;
+expectType<bigint>(test9);
+
+declare const test10: HomomorphicPick<symbol, never>;
+expectType<symbol>(test10);
+
+// Picks all keys, if `KeyType` is `any`
+declare const test11: HomomorphicPick<{readonly a: 1; b?: 2} | {readonly c?: 3}, any>;
+expectType<{readonly a: 1; b?: 2} | {readonly c?: 3}>(test11);
+
+// Picks no keys, if `KeyType` is `never`
+declare const test12: HomomorphicPick<{a: 1; b: 2}, never>;
+expectType<{}>(test12);
+
+// Works with index signatures
+declare const test13: HomomorphicPick<{[k: string]: unknown; a: 1; b: 2}, 'a' | 'b'>;
+expectType<{a: 1; b: 2}>(test13);
+
+// Doesn't pick `number` from a `string` index signature
+// eslint-disable-next-line @typescript-eslint/consistent-indexed-object-style
+declare const test14: HomomorphicPick<{[k: string]: unknown}, number>;
+expectType<{}>(test14);

test-d/set-optional.ts

@@ -32,3 +32,7 @@ expectType<{readonly a?: number; b?: string; c?: boolean}>(variation7);
 // Does nothing, if `Keys` is `never`.
 declare const variation8: SetOptional<{a?: number; readonly b?: string; readonly c: boolean}, never>;
 expectType<{a?: number; readonly b?: string; readonly c: boolean}>(variation8);
+
+// Works with index signatures
+declare const variation9: SetOptional<{[k: string]: unknown; a: number; b?: string}, 'a' | 'b'>;
+expectType<{[k: string]: unknown; a?: number; b?: string}>(variation9);

test-d/set-readonly.ts

@@ -32,3 +32,7 @@ expectType<{readonly a?: number; readonly b: string; readonly c: boolean}>(varia
 // Does nothing, if `Keys` is `never`.
 declare const variation8: SetReadonly<{a: number; readonly b: string; readonly c: boolean}, never>;
 expectType<{a: number; readonly b: string; readonly c: boolean}>(variation8);
+
+// Works with index signatures
+declare const variation9: SetReadonly<{[k: string]: unknown; a: number; readonly b: string}, 'a' | 'b'>;
+expectType<{[k: string]: unknown; readonly a: number; readonly b: string}>(variation9);

test-d/set-required.ts

@@ -36,3 +36,7 @@ expectType<{readonly a: number; b: string; c: boolean}>(variation8);
 // Does nothing, if `Keys` is `never`.
 declare const variation9: SetRequired<{a?: number; readonly b?: string; readonly c: boolean}, never>;
 expectType<{a?: number; readonly b?: string; readonly c: boolean}>(variation9);
+
+// Works with index signatures
+declare const variation10: SetRequired<{[k: string]: unknown; a?: number; b: string}, 'a' | 'b'>;
+expectType<{[k: string]: unknown; a: number; b: string}>(variation10);