Pick, Omit and union types in TypeScript

TypeScript’s Built in Pick, Omit and other similar type operators don’t quite support union types well.

One of the main issues is use of keyof in there definitions:

type Pick<T, K extends keyof T> = {
  [P in K]: T[P];
};

type Omit<T, K extends keyof any> = Pick<T, Exclude<keyof T, K>>;

By default when keyof is used with union type it will count only common keys. This is usually not what you want, for example:

type X =
  | { type: "string"; foo: string; komara: 1 }
  | { type: "number"; foo: number; kzxzdad: string }
  | { type: "undefined"; foo: undefined; ajjs: 1; asdad: 44 };

type Keys<T> = keyof T;

// type Keys_X = "type" | "foo"
type Keys_X = Keys<X>;

To fix this we want to use distributive conditional types:

type DistributiveKeys<T> = T extends unknown ? Keys<T> : never;

// type DistributiveKeys_X = "type" | "foo" | "komara" | "kzxzdad" | "ajjs" | "asdad"
type DistributiveKeys_X = DistributiveKeys<X>;

Now we have the foundation we need to rebuild distributive Pick and Omit. First attempt might look like this:

type DistributivePick<T, K extends DistributiveKeys<T>> = T extends unknown
  ? Pick<T, Extract<keyof T, K>>
  : never;

export type DistributiveOmit<
  T,
  K extends DistributiveKeys<T>
> = T extends unknown ? Omit<T, Extract<keyof T, K>> : never;

This would work, but it would have very ugly computed type when you are introspecting types in an IDE. I’ve found this version to result in much nicer type:

type DistributivePick<T, K extends DistributiveKeys<T>> = T extends unknown
  ? { [P in keyof Pick_<T, K>]: Pick_<T, K>[P] }
  : never;

type Pick_<T, K> = Pick<T, Extract<keyof T, K>>;

export type DistributiveOmit<
  T,
  K extends DistributiveKeys<T>
> = T extends unknown ? { [P in keyof Omit_<T, K>]: Omit_<T, K>[P] } : never;

type Omit_<T, K> = Omit<T, Extract<keyof T, K>>;

To see this in action we can compare DistributivePick, a naive implementation of DistributivePick and Pick:

// type Test1Pick = {
//     foo: string | number | undefined;
//     type: "string" | "number" | "undefined";
// }
type Test1Pick = Pick<X, "foo" | "type">;

// type Test1 =
// | { type: "string"; foo: string; }
// | { type: "number"; foo: number; }
// | { type: "undefined"; foo: undefined; }
type Test1 = DistributivePick<X, "foo" | "type">;

type Test2Pick = Pick<X, "foo" | "type" | "kzxzdad">;
// ERROR             ~~~~~~~~~~~~~~~~~~~~~~~~~~
// Type '"type" | "foo" | "kzxzdad"' does not satisfy the constraint '"type" | "foo"'.
//   Type '"kzxzdad"' is not assignable to type '"type" | "foo"'.(2344)

// type Test2 =
// | { type: "string"; foo: string; }
// | { type: "number"; foo: number; kzxzdad: string; }
// | { type: "undefined"; foo: undefined; }
type Test2 = DistributivePick<X, "foo" | "type" | "kzxzdad">;

Removing {} from result

Sometimes when using those DistributivePick and DistributiveOmit you might get type that contains ... | {} | ... and in some cases you might want to remove it.

type RemoveObject<T> = T extends unknown
  ? keyof T extends never
    ? never
    : T
  : never;

type Test3 =
  | {}
  | { az: string }
  | { foo: 12 }
  | {
      foo: string;
      __foo: string;
    };

// type Test3 = {
//     az: string;
// } | {
//     foo: 12;
// } | {
//     foo: string;
//     __foo: string;
// }
type Test4 = RemoveObject<Test3>;

Now let’s apply this technique to DistributivePick and DistributiveOmit

export type DistributivePick<
  T,
  K extends DistributiveKeys<T>
> = T extends unknown
  ? keyof Pick_<T, K> extends never
    ? never
    : { [P in keyof Pick_<T, K>]: Pick_<T, K>[P] }
  : never;

type Pick_<T, K> = Pick<T, Extract<keyof T, K>>;

export type DistributiveOmit<
  T,
  K extends DistributiveKeys<T>
> = T extends unknown
  ? keyof Omit_<T, K> extends never
    ? never
    : { [P in keyof Omit_<T, K>]: Omit_<T, K>[P] }
  : never;

type Omit_<T, K> = Omit<T, Extract<keyof T, K>>;

// type Test5 = {
//     az: string;
// } | {
//     __foo: string;
// }
type Test6 = DistributiveOmit<Test3, "foo">;

// type Test5 = {
//     foo: 12;
// } | {
//     foo: string;
// }
type Test5 = DistributivePick<Test3, "foo">;

Originally commented about this here.