🐣

Union typesに対してPickやOmitをしたい時

に公開
TypeScript
tech

タイトルの通りです。
普通（？）のオブジェクト型に対して一部のプロパティを取り出したり削除したりしたい時にはPickやOmitを使いますね。
ただPickとOmitはUnion typesに対して分配的には効きません（後述）。
そこでUnion typesに対して分配的にPickやOmitをする方法を考えていきます。
実際に作ろうと思ったときと同じ手順で説明していくので、これからGeneric typesを書いていきたい方の参考にもなればと思います。

結論

// 補助的な型
type KeyOfUnion<T> = T extends T ? keyof T : never;

// 分配的なPick
type DistributivePick<T, K extends KeyOfUnion<T>> = T extends T
  ? Pick<T, Extract<K, keyof T>>
  : never;

// 分配的なOmit
type DistributiveOmit<T, K extends KeyOfUnion<T>> = T extends T
  ? Omit<T, K>
  : never;

通常のPick, Omitではダメなのか

次のようなT1からT2を作りたいとします。
つまりUnion type T1に対して、Unionの構造を保ったままb, cを取り出す、もしくはaを削除したいです。

type T1 =
  | {
      a: string;
      b: true;
      c: number;
    }
  | {
      a: string;
      b: false;
    };
type T2 =
  | {
      b: true;
      c: number;
    }
  | {
      b: false;
    };

ここで素直にPick, Omitを使おうとしてみます。

type S = Pick<T1, "b" | "c">; // error
type U = Omit<T1, "a">; // = { b: boolean; }

しかしうまくいきません。
Pickを使った方はcなんてキーはT1には無いよと怒られてしまいます。
Omitを使った方はエラーは出ませんが、結果の型は{ b: boolean; }になってしまっていて意図したものとは異なります。
PickとOmitの定義には触れませんが、Pickで指定できるキーやOmitで取り出せるプロパティは全てのUnionに共通するもののみです。

とりあえず分配してみる

Generic typesを分配したいときにまず思いつくのはConditional Typesです。
https://www.typescriptlang.org/docs/handbook/2/conditional-types.html#distributive-conditional-types

特にT extends T ? Hoge<T> : neverと書くことでTにUnion typesが来た時にUnionを分配してHogeに渡してくれます。
やってみましょう。

type DistributivePick<T, K extends keyof T> = T extends T ? Pick<T, K> : never;
type DistributiveOmit<T, K extends keyof T> = T extends T ? Omit<T, K> : never;

type S = DistributivePick<T1, "b" | "c">; // error
type U = DistributiveOmit<T1, "a">;

const u1: U = {
  b: true,
  c: 1,
};
const u2: U = {
  b: false,
};

DistributivePickは相変わらずエラーが出ていますが、DistributiveOmitは良さそうですね！

キーの制約を直す

ここまでのDistributivePickの定義ではcというキーはないんだが？？？と怒られています。
理由はkeyofもUnion typesに対しては全てのUnionに共通するものしか取れないからです。

type K1 = keyof T1; // = "a" | "b"

ということはここまでの話と同じようにkeyofにも分配して渡してしまえば良さそうです。

type KeyOfUnion<T> = T extends T ? keyof T : never;

type K2 = KeyOfUnion<T1>; // = "a" | "b" | "c"

良さそうですね。
キーの制約をkeyof TからKeyOfUnion<T>に変えてみましょう。

type DistributivePick<T, K extends KeyOfUnion<T>> = T extends T
  ? Pick<T, K>
  : never;
type DistributiveOmit<T, K extends KeyOfUnion<T>> = T extends T
  ? Omit<T, K>
  : never;

type S = DistributivePick<T1, "b" | "c">;

const s1: S = {
  b: true,
  c: 1,
};
const s2: S = { // error
  b: false,
};

DistributiveOmitについてもUnionのある要素にしか存在しないキーを指定したいことがあると思うので、同様の修正をしています。

型定義でエラーは出なくなりました。
ただ今度は型を使っている箇所でエラーが出てしまいました。
最後にこれを直して完成させましょう。

Pickするキーを絞り込む

ここまででDistributivePickを使ったSは実は{ b: true; c: number; } | { b: false; c: unknown; }という型になっています。
bがfalseの時にcがunknownになってしまうのが問題です。
これはT1が分配された結果、Pick<{ a: string; b: false; }, "b" | "c">のような型計算が起きているからです。
cは{ a: string; b: false; }には存在しないキーなので本来はPickできません。
そのためunknownになってしまうようです。

Pickで指定するキーを第一型引数に存在するキーのみに絞ればよさそうです。
Union typesに対して絞り込みを行うにはExcludeがあります。
Omitは存在しないキーを渡そうが関係ないので、今度はDistributivePickのみを修正します。

type DistributivePick<T, K extends KeyOfUnion<T>> = T extends T
  ? Pick<T, Exclude<K, keyof T>>
  : never;

type S = DistributivePick<T1, "b" | "c">;

const s1: S = {
  b: true,
  c: 1,
};
const s2: S = {
  b: false,
};

LGTM!

完成

ここまでで書いたコードのまとめです。

type T1 =
  | {
      a: string;
      b: true;
      c: number;
    }
  | {
      a: string;
      b: false;
    };

type KeyOfUnion<T> = T extends T ? keyof T : never;
type DistributivePick<T, K extends KeyOfUnion<T>> = T extends T
  ? Pick<T, Exclude<K, keyof T>>
  : never;
type DistributiveOmit<T, K extends KeyOfUnion<T>> = T extends T
  ? Omit<T, K>
  : never;

type S = DistributivePick<T1, "b" | "c">;
const s1: S = {
  b: true,
  c: 1,
};
const s2: S = {
  b: false,
};

type U = DistributiveOmit<T1, "a">;
const u1: U = {
  b: true,
  c: 1,
};
const u2: U = {
  b: false,
};

皆様も良きTSライフを！

Discussion