Zenn
🤖

Pythonのzip関数の罠を回避する

2024/02/17に公開
Python
tech

はじめに

Pythonのzip関数は、複数のイテラブルを並列に反復処理するための強力なツールである。たとえば、複数のリストを同時にループさせる処理を、簡潔で読みやすく書ける。

a = [1, 2, 3]
b = ["a", "b", "c"]
for x, y in zip(a, b):
    print(x, y)

このコードは、次を出力する。

1 a
2 b
3 c

しかし、異なる長さのイテラブルに対するzip関数の処理は、サイレントエラーにつながる可能性がある。zip関数は、警告なしにもっとも短いイテラブルに対して処理を終了してしまう。そのため、ソフトウェア開発においてバグを見過ごす可能性がある。

zip関数の罠

核となる問題は、zip関数を異なる長さの複数のイテラブルに対して使用した場合に発生する。仕様上、zipは最短のイテラブルを使い切った時点で、反復処理を終了する。そして、この動作では、処理されなかった要素が存在するという事実をプログラマに警告しない。

この挙動はバグの原因となりうる。たとえば、次のコードを考える。

a = [1, 2, 3]
b = ['a', 'b', 'c']
a.append(4)
for x, y in zip(a, b):
    print(x, y)

このコードは、次を出力する。

1 a
2 b
3 c

このコードでは、aに新要素を追加したものの、bには新要素を追加し忘れている。しかし、zip関数は最短のイテラブルを使い切った時点で処理を終了してしまうため、長さの短いbに合わせて処理してしまい、aに追加された新要素は出力されない。そのうえ、長さが異なることに関して何も警告を出力しないのだ。

zip関数の罠を回避する

このようなzip関数の罠を回避する方法は2つある。

  1. 同じ長さを保証する
  2. 長さが長いほうに合わせる

どちらが優れているというのはない。どちらを利用すべきかは、要件により異なる。

同じ長さを保証する

この罠を避けるもっとも簡単な方法は、zip関数を使う前に、すべてのイテラブルが同じ長さであることを確認することである。たとえば、次のように手動で確認できる。

a = [1, 2, 3]
b = ["a", "b", "c"]
a.append(4)
if len(a) != len(b):
    raise ValueError("length of a and b must be the same")
for x, y in zip(a, b):
    print(x, y)

Python 3.10からは zip関数にstrictオプションが追加された。strictオプションをTrueにすると、同じ長さのイテラブルを強制し、長さが異なる場合はValueErrorを発生させられる。たとえば、次のように使う。

a = [1, 2, 3]
b = ["a", "b", "c"]
a.append(4)
for x, y in zip(a, b, strict=True):
    print(x, y)

このコードは、次のエラーを発生させる。

ValueError: zip() argument 2 is shorter than argument 1

このように、複数のイテラブルの長さが同じであることを保証すれば、zip関数の罠を回避できる。

長さが長いほうに合わせる

異なる長さのイテラブルが想定される場合、zip関数の代わりに、Pythonの組み込みモジュールitertoolsにある、itertools.zip_longest関数を使える。この関数は、長さが異なるイテラブルを渡したとき、長さが短いイテラブルの欠損値をfillvalueに指定した値[1]で埋め、すべての要素を処理できるようにする。たとえば、次のように使う。

import itertools

a = [1, 2, 3]
b = ["a", "b", "c"]
a.append(4)
for x, y in itertools.zip_longest(a, b, fillvalue="N/A"):
    print(x, y)

このコードは、次を出力する。

1 a
2 b
3 c
4 N/A

この実装では、aに追加された新要素が出力され、新要素を追加し忘れたbにはfillvalueで指定した値（今回は"N/A"）が出力される。このように、zip_longest関数を使うことで、長さが異なることに気づける。

まとめ

Pythonのzip関数は紛れもなく便利である。しかし、長さが等しくないイテラブルを処理すると、気づきにくいバグを引き起こす可能性がある。同じ長さを保証するか、長さの違いを明示的に処理すれば、より堅牢でエラーを意識したコードにできる。

参考

脚注

  1. なお、fillvalueのデフォルト値はNoneである。 ↩︎

Discussion