TypeScriptでタグ付きユニオンのタグがユニオンの場合でも個別にExtractしたい

何を言っているんだ?というタイトルだけど、今日はTypeScriptの型で遊んでた。

タグ付きユニオンから型をExtractしたい

こういうタグ付きユニオンがあって

type MyUnion =
    | { tag: "a", value: boolean }
    | { tag: "b", value: boolean }
    | { tag: "c", value: string }

その中の型をExtractしたいときは、こんな風に書ける

type A = Extract<MyUnion, { tag: "a" }>

そうするとtagが"a"の型を取得できるので、このテストがとおる

import type { Equal, Expect } from '@type-challenges/utils'

type cases = [
    Expect<Equal<A, { tag: "a", value: boolean }>>,
]

タグがユニオンのときはとおらない

ここで「"a"と"b"のタグのvalueの型はどっちもbooleanだから、tagの方をこんな風にunionにしてしまおう!」ってすると

type MyUnion =
    | { tag: "a" | "b", value: boolean }
    | { tag: "c", value: string }

同じようにExtractしようとしても

type A = Extract<MyUnion, { tag: "a" }>

neverになってしまう(このテストがとおる)

type cases = [
    Expect<Equal<A, never>>,
]

"a" | "b"でExtractすれば取れるけど、抽出された型のタグも"a" | "b"になっている。

type A = Extract<MyUnion, { tag: "a" | "b" }>

type cases = [
    Expect<Equal<A, { tag: "a" | "b", value: boolean }>>,
]

なるほど。

こういうことをしたい

タグがユニオンの場合でも個別にExtractしてみたいのだ。↓みたいな感じ。でも、これは今見たとおりAとBがneverになるのでとおらない。

type MyUnion =
    | { tag: "a" | "b", value: boolean }
    | { tag: "c", value: string }

type A = Extract<MyUnion, { tag: "a" }>
type B = Extract<MyUnion, { tag: "b" }>
type C = Extract<MyUnion, { tag: "c" }>

type cases = [
    Expect<Equal<A, { tag: "a", value: boolean }>>,
    Expect<Equal<B, { tag: "b", value: boolean }>>,
    Expect<Equal<C, { tag: "c", value: string }>>,
]

どうしようかな?

MyUnionのtagを最初みたいに個別で書けばとおるんだから、unionで定義されたタグを個別のタグに分解できればやりたいことができそう。と思ってごにょごにょやってみた。

まずは、MyUnionをMapped Typesに変換する

type MyUnion =
    | { tag: "a" | "b", value: boolean }
    | { tag: "c", value: string }

type MyMap = { [I in MyUnion as I["tag"]]: I["value"] }

// MyMapは、こうなってる
type cases = [
    Expect<Equal<MyMap, { a: boolean, b: boolean, c: string }>>,
]

KeyのRemappingを使って、tagをキーにvalueをその型にしている。(参照: https://www.typescriptlang.org/docs/handbook/2/mapped-types.html#key-remapping-via-as

このときキーがunionの場合は分解してくれるからaとbが分解される。

次に、このMapped Typesからタグ付きユニオンを作り直す。

type MyUnion2 = {
    [K in keyof MyMap]: { tag: K, value: MyMap[K] }
}[keyof MyMap]

// MyUnion2は、こうなってる
type cases = [
    Expect<Equal<MyUnion2,
        | { tag: "a", value: boolean }
        | { tag: "b", value: boolean }
        | { tag: "c", value: string }>>,
]

全然この件とは関係ないけど、型について調べてるときに見かけたこの Issue を参考にして変換を書いた。

Fix multiple issues with indexed access types applied to mapped types by ahejlsberg · Pull Request #47109 · microsoft/TypeScript · GitHub

ここにこういうコードがあって「へーそんなやり方があるのかー」って思ってたから「あ、ここで使えそう」ってなったのだ。

type UnionRecord = { [P in keyof RecordMap]: RecordType<P> }[keyof RecordMap];

この変換に対する自分の理解を説明すると

type Foo = {
    [K in keyof MyMap]: { tag: K, value: MyMap[K] }
}

こうするとMyMapのキーに対して、こんなMapを新たに作り出す

type Foo = {
    a: {
        tag: "a";
        value: boolean;
    };
    b: {
        tag: "b";
        value: boolean;
    };
    c: {
        tag: "c";
        value: string;
    };
}

このMapに対して最後に [keyof MyMap] をつけて、値の部分を取り出している。という仕組み。すごいなぁ。

ということで

こうなった。わーい。

type MyUnion =
    | { tag: "a" | "b", value: boolean }
    | { tag: "c", value: string }

type MyMap = { [I in MyUnion as I["tag"]]: I["value"] }

type MyUnion2 = {
    [K in keyof MyMap]: { tag: K, value: MyMap[K] }
}[keyof MyMap]

type A = Extract<MyUnion2, { tag: "a" }>
type B = Extract<MyUnion2, { tag: "b" }>
type C = Extract<MyUnion2, { tag: "c" }>

/* _____________ TestCase _____________ */

import type { Equal, Expect } from '@type-challenges/utils'

type cases = [
    Expect<Equal<A, { tag: "a", value: boolean }>>,
    Expect<Equal<B, { tag: "b", value: boolean }>>,
    Expect<Equal<C, { tag: "c", value: string }>>,
]

ってことで、型で遊んで楽しかった。

もっと簡単なやり方ないかな?

僕は↑のやり方しかできなかったんだけど、もっと簡単なやり方があったらいいなぁと思った。知ってる人はTwitterとかで教えてくれると嬉しい!