【TypeScript】zip 関数を作成して型をつける

Python や Ruby には組み込みで zip 関数（メソッド）があります。

例えば、Python の zip 関数は以下のように使います（[公式リファレンスから引用](https://docs.python.org/ja/3.12/library/functions.html#zip)）。

```py:Python の zip 関数(引用)
list(zip(range(3), ['fee', 'fi', 'fo', 'fum']))
# [(0, 'fee'), (1, 'fi'), (2, 'fo')]
# ※注: Python の zip 関数はタプルのイテレータを返すので、リストを構築するのに list() を使用
```

Python の zip 関数では、引数の数は可変長で、リストだけでなくイテラブルなら受け付けるようになっています。

あれば便利かなと思う関数ですが、JavaScript には組み込みでこのような関数・メソッドはありません。
このような関数を自前で実装するのは、シンプルな実装であればそこまで難しくありませんし、実装例も探せばそれなりに見つかります。ただ、TypeScript で使うために汎用的な型をつけようとしている記事はあまりみかけません。

ということで、この記事では TypeScript の機能を活用して可変長引数を取る zip 関数にそれなりに正確な型をつけるのを目標とします。
TypeScript のバージョンは **5.1.6** で、`--strict` オプションを有効にしています。

## TL;DR

- イテレータを引数としてジェネレータを返す場合

```ts
function* zip<T extends Iterable<any>[]>(
  ...iterables: [...T]
): Generator<{ [K in keyof T]: T[K] extends Iterable<infer U> ? U : never }, void, unknown> {
  const iterators = iterables.map((it) => it[Symbol.iterator]());
  while (iterators.length) {
    const result = [];
    for (const it of iterators) {
      const elemObj = it.next();
      if (elemObj.done) {
        return;
      }
      result.push(elemObj.value);
    }
    yield result as any;
  }
}
```

- 配列を引数として配列を返す場合

```ts
const zipArray = <T extends (readonly unknown[])[]>(
  ...args: [...T]
): { [K in keyof T]: T[K][number] }[] => {
  if (!args.length) return [];
  const minLen = args.reduce((a, c) => (a.length < c.length ? a : c)).length;
  let result = [];
  for (let i = 0; i < minLen; i++) {
    result.push(args.map((arg) => arg[i]));
  }
  return result as any[];
};
```

## zip 関数の定義

まず、JavaScript の標準にない zip 関数をどう定義するかです。

先ほど書いた通り、Python の zip 関数は、イテラブル（リスト等）を可変長で引数にとり、イテレータを返します。
もう1つ重要な点として、zip で返されるイテレータは、引数の中で最も短いイテラブルを消費しきった時点でイテレーションを終了します。上の Python の例でいえば、`'fum'` は返されたイテレータからは出てきません。

今回はこの Python での zip 関数を真似して実装していきます。JavaScript に合わせて言語化するならこんな感じでしょうか。

- zip は関数で、イテラブルを可変長引数に取る。
- 引数の要素を一つずつ取り出して新たなタプルのイテレータを生成する。
- 生成したイテレータは、引数のイテラブルの中で最も短いイテラブルが尽きた時点でイテレーションを止める。

ここでいうタプルは、TypeScript でのタプル型なので、JavaScript では普通の配列です。

イテレータのほうが汎用性は高いのですが、イテレータを配列に直すにはひと手間必要なのと、実用上配列だけ扱えれば十分な場面が多いかもしれません。
そのため、引数・返り値ともに配列となるような zip 関数も作ってみたいと思います。

<details><summary>（補足）イテラブル、イテレータ</summary>

<br>

用語が似ててややこしいので、少しだけ補足です。
詳しい説明は MDN の[反復処理プロトコル](https://developer.mozilla.org/ja/docs/Web/JavaScript/Guide/Iterators_and_Generators)にあります。[JavaScriptのIterator / Generatorの整理](https://zenn.dev/qnighy/articles/112af47edfda96)などの記事も参考になります。

要点だけ言うと

- **イテラブル**は、`[Symbol.iterator]()` メソッドでイテレータを返すインターフェース
- **イテレータ**は、`next()` メソッドで done や value プロパティを持つオブジェクトを返すインターフェース

スプレッド構文で展開できたり、`for...of` 構文で反復処理できるのは**イテラブル**オブジェクトです。配列に対してスプレッド構文を適用できるのは、イテラブル、つまり `[Symbol.iterator]()` が実装されているからです。一方、配列には `next()` メソッドはないためイテレータではありません。

イテラブル、イテレータの関係は `for...of` 構文を `while` で書き換えてみると分かりやすいかもしれません。イテラブルの `[Symbol.iterator]()` を呼び出してイテレータを生成し、それを使って順番に処理していくイメージです。

```ts: for...of 構文
const arr = [1, 2, 3];
for (const value of arr) {
  // ...
}
```

```ts: for...of 構文を while に書き換え
const arr = [1, 2, 3];
const iter = arr[Symbol.iterator](); // iter の型は IterableIterator<number>
while (true) {
  const { value, done } = iter.next();
  if (done) break;
  // ...
}
```

なお、イテレータオブジェクトの多くはイテラブルも実装しており、今回の配列から生成した `iter` もイテレータかつイテラブルです。
そのようなインターフェースは、TypeScript では `IterableIterator` で表現されます。
`iter` はイテラブルなので、`for...of` 構文で使用することもできます。

```ts: iter を for...of 構文で使用
const arr = [1, 2, 3];
const iterableiter = arr[Symbol.iterator]();
// OK
for (const value of iterableiter) {
  // ...
}

// 配列を Iterable<number> にアップキャスト
const iterable: Iterable<number> = [1, 2, 3];
const iter_n = iterable[Symbol.iterator](); // Iterator<number>
// 型エラー！(イテラブルでないイテレータの場合はfor...ofを使えない)
for (const value of iter_n) {
  // ...
}
```

<hr>
</details>

## zip 関数の実装

JavaScript でイテレータを生成するならジェネレータ関数を使うのが簡単で確実です（ジェネレータオブジェクトは Iterator, Iterable を実装しています）。
先ほどの定義を満たす zip 関数は __sil さんの記事 [JavaScriptでPythonのzip関数を実装してみました](https://qiita.com/__sil/items/d7a83d4072ae47ad404c) で作成されていますので、このままお借りしました。記事の中で詳しい解説もされていますので、ここでの実装の説明は省かせていただきます。

```ts: ジェネレータを返す zip 関数（引用、一部改変）
function* zip(...iterables) {
  const iterators = iterables.map((it) => it[Symbol.iterator]());
  while (iterators.length) {
    const result = [];
    for (const it of iterators) {
      const elemObj = it.next();
      if (elemObj.done) {
        return;
      }
      result.push(elemObj.value);
    }
    yield result;
  }
}
```

この関数を使うとこんな感じです。配列でなくあえて `Set` を入れてみています。

```ts
for (const el of zip(new Set([1, 2]), ["a", "b"])) {
  console.log(el); // [1, "a"] と [2, "b"] を出力
}
```

この時点では zip 関数に全く型をつけていませんが、TypeScript の型推論では引数 `iterables` の型は `any[]`、返り値の型は `Generator<any[], void, unknown>` となっています。

引数 `iterables` の型が配列なのは、TypeScript において可変長引数の型は配列で宣言することができるからです。この場合各々の引数が `any` 型である可変長引数を意味します。  
今回、引数に取れるのはイテラブルです。なので `any` を TypeScript でイテラブルインターフェースを表す `Iterable<any>` に置き換えます。なお、`Iterable` は [lib.es2015.iterable.d.ts](https://github.com/microsoft/TypeScript/blob/v5.1.6/lib/lib.es2015.iterable.d.ts#L41-L50) で定義されています。

返り値の型の `Gererator` はあまり見慣れない型かもしれません。ただ、この記事の範囲内では第1型引数の型が `for...of` やスプレッド構文で取り出せる型と思ってもらえば支障はないです。例えば、先ほどの例の `el` の型は `Generator<any[], void, unknown>` の第1型引数と同じ `any[]` となっています。

ここまででとりあえず型をつけてみると以下のようになります。

```ts
function* zip(...iterables: Iterable<any>[]): Generator<any[], void, unknown> {
  // 省略
}
```

これで TypeScript の型検査は通るようになりましたが、このままだと返り値の型に `any` が出てきて非常に使いにくいです。
これをなんとかしようというのが今回の目標です。

それと、配列を引数に取って配列を返す `zip` 関数も定義してみます。
色々書き方はあると思いますが、今回は素朴にこのように定義してみます。

```ts: 配列を返す zip 関数
const zipArray = (...args: any[][]): any[][] => {
  if (!args.length) return []
  const minLen = args.reduce((a, c) => a.length < c.length ? a : c).length
  let result = []
  for (let i = 0; i < minLen; i++) {
    result.push(args.map(arg => arg[i]))
  }
  return result
}
```

## ジェネリクスを利用

では早速型をつけていきたいと思いますが、その前にまずは引数が固定長の場合どうなるかというのを考えてみます。  
TypeScript で引数の型を返り値に使いたい場合に、まず最初に思いつくのがジェネリクスを使う方法かと思います。引数が2つの場合の `zip2` 関数の型は以下のようになります。`zip2` 関数内部での型エラーはひとまず無視です。

```ts: 引数が2個の zip 関数
function* zip2<T1, T2>(
  iter1: Iterable<T1>,
  iter2: Iterable<T2>,
): Generator<[T1, T2], void, unknown> {
  const iterators = [iter1, iter2].map((it) => it[Symbol.iterator]());
  // 省略
}
// el の型は [number, string | number]
for (const el of zip2(new Set([1, 2]), ["a", "b", 3])) {
  console.log(el); // [1, "a"] と [2, "b"] を出力
}
```

ここで、2つの引数をもつ関数は、可変長引数の型を長さ2のタプルにすることでも表現できます。つまり、先ほどの `zip2` は以下のようにも書けます。

```ts: 引数が2個の zip 関数（引数にタプルを使用）
function* zip2<T1, T2>(
  ...iterables: [Iterable<T1>, Iterable<T2>]
): Generator<[T1, T2], void, unknown> {
  const iterators = iterables.map((it) => it[Symbol.iterator]());
  // 省略
}
```

こうする方が可変長引数の場合と形式を揃えられるので、以後引数はこのように書いていきます。

`el` の型は `[number, string | number]` となっており、いい感じに型がつけられているのが分かります。ただし、当然ながらこの `zip2` 関数は引数が2個のときにしか使えません。引数の数が3個、4個、…の場合は、型推論のために引数の数に対応した zip 関数を量産する必要があります。それは嫌です。

こんなことをせずとも、幸い TypeScript には関数の型をオーバーロードする機能があります。これを使うと、可変長の引数の関数に対して、引数の数を固定した場合の型を指定することができます。さらに、可変長引数に対してもジェネリクスを使うことはできるので、`any` の代わりにジェネリクスを利用した関数の型を宣言できます。

具体的には次のようになります。

```ts: オーバーロードを使った関数の型定義
function zip<T1>(...iterables: [Iterable<T1>]): Generator<[T1], void, unknown>;
function zip<T1, T2>(...iterables: [Iterable<T1>, Iterable<T2>]): Generator<[T1, T2], void, unknown>;
function zip<T1, T2, T3>(...iterables: [Iterable<T1>, Iterable<T2>, Iterable<T3>]): Generator<[T1, T2, T3], void, unknown>;
function zip<T>(...iterables: Iterable<T>[]): Generator<T[], void, unknown>;
function zip(...iterables: Iterable<any>[]): Generator<any[], void, unknown>;
function* zip(...iterables: Iterable<any>[]): Generator<any[], void, unknown> {
  // 省略
}

// el の型は [number, string | number, string]
for (const el of zip(new Set([1, 2]), ["a", "b", 3], ["c"])) {
  console.log(el);
}
// el の型は (string | number)[]
for (const el of zip(new Set([1, 2]), ["a", "b", 3], ["c"], [4, 5, 6])) {
  console.log(el);
}
// el の型は any[]
for (const el of zip(new Set([1, 2]), ["a", "b"], ["c"], [4, 5, 6])) {
  console.log(el);
}
```

引数が1個のときは一番上の型、2個のときは2番目、3個のときは3番目の型が使われるのでタプルに推論されます。それ以上あっても可能であれば `T[]` を返すようにして、それ以外は `any` を使う型にフォールバックします。  
たくさん書けばそれだけ引数の数が多くなっても対応できます。ただ、オーバーロードを無限に書いていくわけにはいかないので、対応できる数に限りは出てきてしまいます。とはいえ、実用上は、特殊な場合を除き多くとも5～6個あれば十分かと思います。

結構な力業ですが、似たようなアプローチを取っているケースは有名ライブラリ（例えば[@types/node](https://github.com/DefinitelyTyped/DefinitelyTyped/blob/3cd3393e521f755e923751821b08dd6264431143/types/node/util.d.ts#L972-L999)）でも意外とあります。

なお、`el` の型が `any[]` となってしまうのを回避するのに、`zip<string | number>(new Set([1, 2]), ["a", "b"], ["c"], [4, 5, 6])` のように明示的に型引数を指定してあげることもできます。  
ですが、いちいち型引数を指定するのは面倒ですし、その場合でも `el` がタプルでなく配列に推論されてしまうため、使いにくさは残ります。

## Mapped Types

ここまでジェネリクスを使って色々試してみましたが、どれも完全にはうまくいきませんでした。  
実現したい関数の型は以下のようなものです。

```ts: こんな感じに実現したいというイメージ（TypeScript による解釈はできません）
function zip<T1, T2, ...>(
  ...iterables: [Iterable<T1>, Iterable<T2>, ...]
): Generator<[T1, T2, ...], void, unknown>;
```

引数と返り値のタプル型の部分にだけ着目すると、`[Iterable<T1>, Iterable<T2>, ...]` のようなタプル型から `[T1, T2, ...]` というタプル型に変換できれば解決できそうです。

天下り的ですが、引数のタプルの型を返り値で使うタプルの型に変換する、というのは実は `Promise.all` の型定義でも行われています。`Promise.all` では、要素が `Promise` である配列を引数に取り、非同期に履行(fulfilled)された値の配列を `Promise` でくるんだものを返します。

```ts
// 引数の型が values: [Promise<number>, Promise<string>]、返り値の型が Promise<[number, string]>
const promises = Promise.all([Promise.resolve(42), Promise.resolve("foo")]);
```

今回やりたいこととかなり似ていますので、これを参考にしたいと思います。

TypeScript 5.1.6 時点での `Promise.all` の[型定義](https://github.com/microsoft/TypeScript/blob/v5.1.6/lib/lib.es2015.promise.d.ts#L39)は以下のようになっています。

```ts: Promise のインターフェース（一部抜粋）
interface PromiseConstructor {
  // ...
  all<T extends readonly unknown[] | []>(values: T): Promise<{ -readonly [P in keyof T]: Awaited<T[P]> }>;
  // ...
}
```

まず、`T extends readonly unknown[] | []` の部分では、`T` が配列かタプルでなければならないという制約を表しています。末尾の `| []` は、引数に配列リテラルが渡されたときに配列でなくタプルに推論されるようにするもののようです。これは制約に `| []` と書く代わりに、引数で `T` を `[...T]` にしても同じなようです[^variadic_tuple_types]。  
返り値の `Promise` の型引数は `{ -readonly [P in keyof T]: Awaited<T[P]> }` となっています。
このような `{ [P in keyof T]: U }` のような型は [Mapped Types](https://www.typescriptlang.org/docs/handbook/2/mapped-types.html) と呼ばれます。
Mapped Types については[こちら](https://zenn.dev/qnighy/articles/dde3d980b5e386)の記事がとてもよくまとまっているのでそちらを見ていただければ事足りますが、今回必要なものだけざっくり説明していきます。

[^variadic_tuple_types]: 型変数に対して `...T` のように書けるのは Variadic Tuple Types という機能のようです。

まず、Mapped Types は主に既存のオブジェクトの型のキーを使って、新しいオブジェクトの型を作るために利用されます。一つだけ例を挙げると、TypeScript の標準ライブラリで提供されている `Pick` という型はこれを使っています。

```ts
type Pick<T, K extends keyof T> = { [P in K]: T[P] };
type PickedObj = Pick<{ a: number; b: string }, "b">; // { b: string }
```

さらに `keyof T` の `T` が型変数であった場合は、Homomorphic Mapped Types という特殊な挙動をします。この特殊な挙動というのはいくつかありますが、このうち、型変数がタプルであった場合は結果がタプルにマップされるという挙動に着目します。例えば、以下の例での `BooleanMap<[string, number]>` は長さ2の `boolean` 型のタプルになります。

```ts
type BooleanMap<T extends unknown[]> = { [K in keyof T]: boolean };
type Mapped = BooleanMap<[string, number]>; // [boolean, boolean]
```

このように、Homomorphic Mapped Types で `{ [P in keyof T]: U }` という形の `U` を変更すれば、配列要素の型を好きなようにマップできます。

`Promise.all` の話に戻りますと、`{ -readonly [P in keyof T]: Awaited<T[P]> }` は `T` が型変数ですので Homomorphic Mapped Type に該当します。さらに、`T` は配列かタプルであるという制約(`T extends readonly unknown[] | []`)があるので、先ほどの話の通り、配列要素が `Awaited<T[P]>` という型にマップされます。  
`T[P]` というのは元の配列要素を表し、`Awaited` は、ここでは中身まで触れませんが、型引数の `Promise` を全て剥がして値を取り出す役割をします。

## Mapped Types を使って定義

ここまででみてきたことを使って、可変長引数の zip 関数を改良してみます。  
まずは、`Promise.all` と同様に、可変長引数の `Iterable<T>[]` の部分を `T` に置き換えて、`T` に制約をつけます。あまり意味がないような気もしますが、一応可変長引数の型は `[...T]` としておきます。一旦、返り値の `T[]` の部分は `Return<T>` と置いておきます。

```ts
function zip<T>(...iterables: Iterable<T>[]): Generator<T[], void, unknown>;
// ↓
function zip<T extends Iterable<unknown>[]>(
  ...iterables: [...T]
): Generator<Return<T>, void, unknown>;
```

次に返り値の型です。これは Homomorphic Mapped Types を利用します。引数タプルのイテラブルをそのままマップするのは `{ [P in keyof T]: T[P] }` のようにすれば可能ですが、ここで欲しいのは Iterable から取り出せる値の型です。  
ひとまず、`Iterable` から値を取り出す役割をする型を `Awaited` に倣って `Nexted` とでも名付けて `zip` 関数の型を記述します。

```ts
function zip<T extends Iterable<unknown>[]>(
  ...iterables: [...T]
): Generator<{ [P in keyof T]: Nexted<T[P]> }, void, unknown>;
```

`Nexted` のような型は [Conditional Types の `infer`](https://www.typescriptlang.org/docs/handbook/2/conditional-types.html#inferring-within-conditional-types) を使って定義できます。公式ドキュメントには `Array` の要素の型を取得する方法が記載されています。

```ts: TypeScript 公式ドキュメントより抜粋
type Flatten<Type> = Type extends Array<infer Item> ? Item : Type;
```

この構文では、`infer` のよって `Array` の要素 `Item` をキャプチャします。この場合、`Type` が配列であるなら、`Item` を返し、そうでなければ `Type` を返す、という意味です。  
これに倣って、`Iterable` の値を取り出したいのであれば、以下のように書けます。

```ts
type Nexted<T> = T extends Iterable<infer U> ? U : never;
```

`T[P]` は各引数の型なので、`Nexted` の型引数に入れるのは必ず `Iterable` となり、偽にはならないので、偽の場合の型は `never` としています。  
これを利用すると、最終的に `zip` 関数は以下のようになります。

```ts
function* zip<T extends Iterable<unknown>[]>(
  ...iterables: [...T]
): Generator<{ [P in keyof T]: T[P] extends Iterable<infer U> ? U : never }, void, unknown> {
  const iterators = iterables.map((it) => it[Symbol.iterator]());
  while (iterators.length) {
    const result = [];
    for (const it of iterators) {
      const elemObj = it.next();
      if (elemObj.done) {
        return;
      }
      result.push(elemObj.value);
    }
    yield result as any;
  }
}
```

`result as any` は `result as { [P in keyof T]: T[P] extends Iterable<infer U> ? U : never }` と書くこともできますが、あまり意味はないので簡単な `any` にしています。

ジェネレータ関数の方が定義できてしまえば、配列の方もほぼ同様に定義できます。  
なお、イテレータから値を取り出すには Condition Types を使う必要がありましたが、配列から値を取り出すだけなら `T[P][number]` とできるのでより簡単に書けます。

```ts
const zipArray = <T extends unknown[][]>(...args: [...T]): { [P in keyof T]: T[P][number] }[] => {
  // 省略
}
```

ただし、配列の場合は引数で readonly な配列やタプルを受け取った場合に型エラーが発生するという問題があります。これを回避するには、制約に readonly を付与します。readonly を付与したからといって、引数に渡すのが必ず readonly でなければならないというわけではないので、特に不都合が発生するわけではありません。

```ts
const zipArray = <T extends (readonly unknown[])[]>(
  ...args: [...T]
): { [P in keyof T]: T[P][number] }[] => {
  if (!args.length) return [];
  const minLen = args.reduce((a, c) => (a.length < c.length ? a : c)).length;
  let result = [];
  for (let i = 0; i < minLen; i++) {
    result.push(args.map((arg) => arg[i]));
  }
  return result as any[];
};
```

## おまけ

実は、Homomorphic Mapped Types で定義できると知る前に、再帰を使って定義したものを作っていました。

```ts
type Zipped<T extends Iterable<unknown>[]> = T extends [Iterable<infer Head>, ...infer Tail]
  ? [Head, ...Zipped<Extract<Tail, Iterable<unknown>[]>>]
  : [];

function* zip<T extends Iterable<unknown>[]>(
  ...iterables: [...T]
): Generator<Zipped<T>, void, void> {
  // ...
}
```

多分こっちの方が応用は効くんだと思いますが、今回に関しては Homomorphic Mapped Types を利用するのがスマートかなと思っています。

## まとめ

Python で実装されている zip 関数の仕様に近い関数に TypeScript で型をつけてみました。この zip 関数の型を作ってこの記事の元となるメモ書きを作ったのは TypeScript 学習したての1年半以上も前だったのですが、可変長引数、Mapped Types、Condition Types、オーバーロードなど、TypeScript の様々な機能を使ってみることになって、実践的に型をつける練習をするのにはちょうどよかったのかなと思います。  
余談ですが、この記事を概ね書いた後に ChatGPT に TypeScript でこのような仕様の zip 関数の型をつけてと聞いたら、何回かのやりとりで今回のとほぼ同じものを出してきました。ちょっと複雑な型でも聞いたらすぐ答えてくれるし、便利な世の中になりましたね。

## 参考

- [TypeScript: Handbook](https://www.typescriptlang.org/docs/handbook/intro.html)
- [JavaScript リファレンス - JavaScript | MDN](https://developer.mozilla.org/ja/docs/Web/JavaScript/Reference)
- [JavaScriptでPythonのzip関数を実装してみました - Qiita](https://qiita.com/__sil/items/d7a83d4072ae47ad404c)
- [TypeScriptの型入門 - Qiita](https://qiita.com/uhyo/items/e2fdef2d3236b9bfe74a)
- [TypeScriptの型初級 - Qiita](https://qiita.com/uhyo/items/da21e2b3c10c8a03952f)
- [TypeScript 4.0で導入されるVariadic Tuple Typesをさっそく使いこなす - Qiita](https://qiita.com/uhyo/items/7e31bbd93a80ce9cec84)
- [Mapped Typesのあれこれ](https://zenn.dev/qnighy/articles/dde3d980b5e386)
- [TypeScript3.6の新機能 - Qiita](https://qiita.com/rana_kualu/items/92ec7f35e245fdc8bd02)
- [any、またお前か——配列とhomomorphic mapped typeの罠](https://zenn.dev/uhyo/articles/array-homomorphic-mapped-type)