matrk's blog

十中八九Python書いてる

zip

zipをfor文でジェネレータ回してたらちょっと引っかかるところがあったのでメモ。以下はpython2.7.5でのログです。

>>> def gen(num):
...   for i in range(num):
...     yield i
...
>>> itr = gen(10)
>>> for t, s in zip(itr, ('a', 'b')):
...   print(t, s)
...
(0, 'a')
(1, 'b')
>>> next(itr)
3
>>> t
1

なんか2が消えました

長さの違うイテレータをzipで使って引っかかりました。
 
for文のitr部分の最後の値が1であり、次は2が返されるはずなのに3が返されました。
これはzipがでくくった順にイテレータを処理しており、3回目のループでタプルがStopIterationを出す前にitrが進められるために発生してるっぽいです。そしてfor文を抜けたあとのtにも2が入っていない事を見るとzip内のどこかでStopIterationが発生するとfor文のtarget_listへの代入処理がされずにループが切られるようです。
ということでイテレータの値が一つ吹き飛んでしまいました。困る。対策はイテレータを短い順にzipでくくること。

>>> for t, s in zip(('a', 'b'), itr):
...   print(t, s)
...
('a', 0)
('b', 1)
>>> next(itr)
2

これで3回目のループでitrが進められる前にタプルがStopIterationを吐くので、for文を抜けたあとitrの次の値が正常にとれます。
イテレータの長さが予測できない場合は潔くキモい書き方をすると思います。主軸にとりたいイテレータを決めてfor文で回し、for文の頭の方で毎回値を取ります。

>>> for t in ('a', 'b'):
...   s = next(itr)
...   print(t, s)
...
('a', 0)
('b', 1)
>>> next(itr)
2

こんなケース無いと思いますが絶対こんなグロいコード書きたくないです。
 
【おまけ】

>>> for _, i in zip([None, None, None], itr):
...   t = next(itr)
...   print(i, t)
...
(0, 3)
(1, 4)
(2, 5)

自分の中ではforは「ループ一回毎にイテレータを一回進めて値を使用する」というイメージがあり、for内でforが回しているイテレータ(expression_list)を進めると値(target_list)がずれていき

>>> for _, i in zip([None, None, None], itr):
...   t = next(itr)
...   print(i, t)
...
(0, 1)
(2, 3)
(4, 5)

となるものだと予想していましたが、実際はfor内でのitrは既にforによって3回進められた後の状態になっているようです。変な書き方はこういうちょっとした部分でバグを呼ぶので気をつけねばなりますまい。