Pick, Omit and union types in TypeScript

You can play along in typescript playground

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

One of the main issues is in keyof. 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:

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;

Tho I prefer to inline Omit and Pick which results in much nicer type shown in the IDE. Best way that I've found for doing this, is:

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 DistributivePickNaive<T, K extends keyof T> = T extends unknown
  ? Pick<T, K>
  : never;

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

// type Test1Naive =
// | { type: "string"; foo: string; }
// | { type: "number"; foo: number; }
// | { type: "undefined"; foo: undefined; }
type Test1Naive = DistributivePickNaive<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 Test2Naive = DistributivePickNaive<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 = DistributivePick2<Test3, "foo">;

Originally commented about this here.