qtatsuの週報

初心者ですわぁ

Pythonの組み込み関数zipに渡したイテレータはコピーされず、元のイテレータが消費される

前書き

タイトル通りの小ネタです。

この挙動を積極的に利用しているコードを散見するのですが、見るたびに混乱するので自分への戒めとしてまとめておきました。

参考リンク

Python公式zip()関数の説明

Built-in Functions — Python 3.9.1 documentation

環境

バージョン
MacOS Catalina 10.15.6
Python3 3.9.0

本題: 何が起こるのか?

タイトルの通り、「zipに渡したイテレータはコピーされず、元のイテレータが消費される」例を示します。

まずはイテレータを作成します。printすると、list_iterator objectがそれぞれ作成されたことを確認できます。

# イテレータの作成と確認
>>> str_it = iter(['a', 'b', 'c'])
>>> int_it = iter([1, 2, 3])

>>> print(str_it, int_it)
<list_iterator object at 0x114320910> <list_iterator object at 0x113e09d60>

これらのイテレータをzip関数に渡します。 zip関数自体も、イテレータを返します。nextを使って値を取り出してみると、zipに渡した二つのイテレータのそれぞれの最初の要素がタプルの形で取り出されていることがわかります。

>>> zipped = zip(str_it, int_it)

>>> next(zipped)
('a', 1)

このタイミングで、zip()関数に渡したstr_itint_itをnextしてみます。 すると、これらのイテレータを直接nextしたのはこれが初めてなのに、両方とも2番目の値が取得されました。

# 最初の要素 'a' ではなく、次の要素 'b' が返ってきた!
>>> next(str_it)
'b'

# 最初の要素 1 ではなく、次の要素 2 が返ってきた!
>>> next(int_it)
2

このままもう一度zippedをnextしてみましょう。 予想通り、先ほどstr_itint_itを一つ進めた影響が反映されているので、イテレータの最後の値が出力されます。

>>> next(zipped)
('c', 3)

zipにイテレータを渡した際の挙動を理解する

これまでにみてきた挙動は、単純に「渡したイテレータオブジェクトが、そのまま消費されている」だけなのですが、ドキュメントを見るとちょっと面白い事実がわかります。

Built-in Functions — Python 3.9.1 documentation

上記リンクから、zipの実装のイメージコード引用します(実際にはCで書かれているはずなので、あくまでイメージですが..)。

def zip(*iterables):
    # zip('ABCD', 'xy') --> Ax By
    sentinel = object()
    iterators = [iter(it) for it in iterables]
    while iterators:
        result = []
        for it in iterators:
            elem = next(it, sentinel)
            if elem is sentinel:
                return
            result.append(elem)
        yield tuple(result)

今回注目したいのは以下の部分です。

 iterators = [iter(it) for it in iterables]

可変長引数として渡されたiterablesについて、組み込みのiter関数を実行しています。 例えば、渡したiterablesがlistだった場合はこのタイミングでイテレータオブジェクトが作成されるだけです。ところが、イテレータが渡された場合に、iter関数は以下のように元のイテレータオブジェクトをそのまま返します。 (下のコード参照)

>>> it = iter([1,2,3])

# イテレータオブジェクトにiter()関数を使う。
>>> it_new = iter(it)

# 全く同じオブジェクト!!
>>> print(f"{id(it)} is {id(it_new)}")
4633907264 is 4633907264

以上の2点から、以下の仕組みがわかりました。