pythonでイテレータを使用するときに、同じイテレータを複数回使用したいときがあります。その時の対処法を２つほどまとめてみました。

目次

問題
解決策
- クラスの組み込み関数__iter__関数を用いる
- itertools.teeで直接コピーする方法
まとめ
参考文献

問題

そもそもイテレータは一回使い切りですので、下記のようにコードを書くと、

it = (i+1 for i in range(10))
print("1回目："+str(list(it)))
print("2回目："+str(list(it)))

出力は以下のようになります。

1回目：[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
2回目：[]

すなわち同じイテレータを二回以上使用することができません。
この時、例外など発生しませんので、注意しないと特定しにくいバグを埋め込んでしまう可能性があります。例えば、簡単な例として、引数にイテレータをとる以下のような関数は注意する必要があります。

def get_all_combinations(X_iterator,Y_iterator):
 return [(X,Y) for X in X_iterator for Y in Y_iterator]

X_iterator = (X for X in range(10))
Y_iterator = (Y for Y in range(10))

print(get_all_combinations(X_iterator,Y_iterator))

実行結果は下のようになり、正しく実行されません。（Y_iteratorを複数回使用できないため、X=0の場合の組み合わせしか得られません。）

[(0, 0), (0, 1), (0, 2), (0, 3), (0, 4), (0, 5), (0, 6), (0, 7), (0, 8), (0, 9)]

解決策

解決策としては「クラスの組み込み関数__iter__関数を用いる」と「itertools.teeを用いる方法」の二つの方法があると思います。

クラスの組み込み関数iter関数を用いる

__iter__関数はforループで呼ばれる関数の一つです。例えば、以下のコードのようにfor文の呼び出し対象としてあるクラスが指定された際に、そのクラスの__iter__関数が呼び出されます。

class SampleClass:
  def __init__(self,N):
    self.N = N

  def __iter__(self):
    for i in range(self.N):
      yield i

Sample = SampleClass(3)
for i in Sample:  # SampleClassの__iter__関数が呼び出されている
  print(i)

出力結果

0
1
2

この__iter__関数を用いることで既存の関数（get_all_combination）を変更しなくても、意図した結果を得ることができます。

X_container = SampleClass(10)
Y_container = SampleClass(10)
print(get_all_combinations(X_container,Y_container))

出力結果を見てみると、確かに全組み合わせを出力できています。

[(0, 0), (0, 1), (0, 2), (0, 3), (0, 4), (0, 5), (0, 6), (0, 7), (0, 8), (0, 9), 
(1, 0), (1, 1), (1, 2), (1, 3), (1, 4), (1, 5), (1, 6), (1, 7), (1, 8), (1, 9), 
(2, 0), (2, 1), (2, 2), (2, 3), (2, 4), (2, 5), (2, 6), (2, 7), (2, 8), (2, 9), 
(3, 0), (3, 1), (3, 2), (3, 3), (3, 4), (3, 5), (3, 6), (3, 7), (3, 8), (3, 9), 
(4, 0), (4, 1), (4, 2), (4, 3), (4, 4), (4, 5), (4, 6), (4, 7), (4, 8), (4, 9), 
(5, 0), (5, 1), (5, 2), (5, 3), (5, 4), (5, 5), (5, 6), (5, 7), (5, 8), (5, 9), 
(6, 0), (6, 1), (6, 2), (6, 3), (6, 4), (6, 5), (6, 6), (6, 7), (6, 8), (6, 9), 
(7, 0), (7, 1), (7, 2), (7, 3), (7, 4), (7, 5), (7, 6), (7, 7), (7, 8), (7, 9), 
(8, 0), (8, 1), (8, 2), (8, 3), (8, 4), (8, 5), (8, 6), (8, 7), (8, 8), (8, 9), 
(9, 0), (9, 1), (9, 2), (9, 3), (9, 4), (9, 5), (9, 6), (9, 7), (9, 8), (9, 9)]

itertools.teeで直接コピーする方法

copy関数ではイテレータをコピーすることができません。そのため、itertools.teeを用います。
例えば初めのコード例をitertools.teeを用いて書くと、以下のようになります。
tee関数の第二引数では複製するiteratorの個数を指定します。

import itertools
it = (i+1 for i in range(10))
it1, it2 = itertools.tee(it,2) #第二引数では複製するiteratorの個数(省略時は2)
print("1回目："+str(list(it1)))
print("2回目："+str(list(it2)))

出力例

1回目：[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
2回目：[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

これを用いて、get_allcombination関数を変更してみます。
やや冗長になりますが、以下のように書けます。

import itertools
def get_all_combinations_update(X_iterator,Y_iterator):
  lst_ = []
  for X in X_iterator:
    Y_iterator,Y_iterator_target = itertools.tee(Y_iterator)
    lst_ += [(X,Y) for Y in Y_iterator_target]
  return lst_

X_iterator = (X for X in range(10))
Y_iterator = (Y for Y in range(10))

print(get_all_combinations_update(X_iterator,Y_iterator))

出力結果を見てみると、こちらも確かに全組み合わせを出力できているとわかります。

[(0, 0), (0, 1), (0, 2), (0, 3), (0, 4), (0, 5), (0, 6), (0, 7), (0, 8), (0, 9), 
(1, 0), (1, 1), (1, 2), (1, 3), (1, 4), (1, 5), (1, 6), (1, 7), (1, 8), (1, 9), 
(2, 0), (2, 1), (2, 2), (2, 3), (2, 4), (2, 5), (2, 6), (2, 7), (2, 8), (2, 9), 
(3, 0), (3, 1), (3, 2), (3, 3), (3, 4), (3, 5), (3, 6), (3, 7), (3, 8), (3, 9), 
(4, 0), (4, 1), (4, 2), (4, 3), (4, 4), (4, 5), (4, 6), (4, 7), (4, 8), (4, 9), 
(5, 0), (5, 1), (5, 2), (5, 3), (5, 4), (5, 5), (5, 6), (5, 7), (5, 8), (5, 9), 
(6, 0), (6, 1), (6, 2), (6, 3), (6, 4), (6, 5), (6, 6), (6, 7), (6, 8), (6, 9), 
(7, 0), (7, 1), (7, 2), (7, 3), (7, 4), (7, 5), (7, 6), (7, 7), (7, 8), (7, 9), 
(8, 0), (8, 1), (8, 2), (8, 3), (8, 4), (8, 5), (8, 6), (8, 7), (8, 8), (8, 9), 
(9, 0), (9, 1), (9, 2), (9, 3), (9, 4), (9, 5), (9, 6), (9, 7), (9, 8), (9, 9)]

まとめ

同じiteratorを複数回使いたいときの方法として。組み込み関数__iter__関数を使用する方法と、itertools.teeを用いる方法の二通りがありました。個人的には、__iter__関数を用いるほうがスマートかなと思いました。

参考文献

Brett Slatkin, "Effective Python", 2016, O'REILLY

補足

イテレータの使用をあきらめる方法

イテレータの使用をあきらめてリスト化する方法をご紹介します。（@shiracamus 様ご指摘ありがとうございました。）

下のコードのように、リスト化してしまえば、多重ループ内でも繰り返して使用することができます。ただし、イテレータでない分、扱う要素数が多くなるとメモリを圧迫しますのでご注意ください。

X_container = [X for X in range(10)]
Y_container = [Y for Y in range(10)]

print(get_all_combinations(X_container,Y_container))