[Tips] Python で複数の Iterator を結合する

Python で複数の Iterator を結合したいことがありました。

やりたいこと

やりたいことは次のようなイメージです。

class SchoolClass:
    def __init__(self, classmates):
        self._classmates = classmates

    def __iter__(self):
        return iter(self._classmates)

classA = SchoolClass(['a', 'b', 'c', 'd'])
classB = SchoolClass(['e', 'f', 'g', 'h'])

# classA のクラスメートを出力
for classmate in classA:
    print(classmate)

# classB のクラスメートを出力
for classmate in classB:
    print(classmate)

# まとめて一気に出力するために、
# classAB = classA と classB を結合した classAB を作りたいけど。。。 
# for classmate in classAB:
#     print(classmate)

解法

ここで、 classA と classB のクラスメイトを個別に出力するのではなく、 Iterator を結合して一気に出力したい場合は、 itertools を使うと簡単です。

import itertools

class SchoolClass:
    def __init__(self, classmates):
        self._classmates = classmates

    def __iter__(self):
        return iter(self._classmates)

classA = SchoolClass(['a', 'b', 'c', 'd'])
classB = SchoolClass(['e', 'f', 'g', 'h'])

# itertools.chain で iterator を結合する
classAB = itertools.chain(classA, classB)

# 一気に出力する
for classmate in classAB:
    print(classmate)

# 出力は
# a
# b
# c
# d
# e
# f
# g
# h

# classB のクラスメートを先に出力することも簡単にできる
classBA = itertools.chain(classB, classA)
for classmate in classBA:
    print(classmate)

# 出力は
# e
# f
# g
# h
# a
# b
# c
# d

itertools.chain() メソッド　で複数の iterator を結合できます。
3つ以上の iterator を結合することもできます。

classABC = itertools.chain(classA, classB, classC)

何が嬉しいか？

今回あげた例では、iterator を結合した上で行う処理が単純な print なので、あまり有り難みがありませんが、例えば、既にある集合に対して iterator で要素を取り出しつつ、各要素に複雑な処理を行なっている既存コードがある際に、その複雑な処理はそのままで処理対象の要素を増やしたい時に便利です。

all_score = 0 # 全員の総合成績を記録する変数
for classmate in classA:
  complex_process1(classmate) # 複雑な処理1
  complex_process2(classmate) # 複雑な処理2
  complex_process3(classmate) # 複雑な処理3
  score = classmate.get_score() # 例えば、成績を返したり
  all_score += score

# 全員の総合成績を出力
print(all_score)

このような処理があったとして、同じ処理を classB に対しても行いたい場合、 iterator の結合を使わないと次のようになると思います。

def calc_all_score(schoolClass):
    all_score = 0
    for classmate in schoolClass:
        complex_process1(classmate) # 複雑な処理1
        complex_process2(classmate) # 複雑な処理2
        complex_process3(classmate) # 複雑な処理3
        score = classmate.get_score() # 例えば、成績を返したり
        all_score += score
    return all_score

all_score = 0 # 全員の総合成績を記録する変数
all_score += calc_all_score(classA)
all_score += calc_all_score(classB)

# 全員の総合成績を出力
print(all_score)

この例ではfor文の処理が短いので関数化してまとめても問題ないのですが、これがとても長い処理だったり、金額計算など正確性が求められる処理である場合、関数化が難しかったり避けたかったりすることがままあると思います。

そのような場合に itertools.chain() を使えば、 iterator を取り扱う側のコードは書き換えずに処理対象を拡大することができます。

import itertools

# 
# 途中のコードは中略
# 

classAB = itertools.chain(classA, classB)

all_score = 0 # 全員の総合成績を記録する変数
for classmate in classAB:
  complex_process1(classmate) # 複雑な処理1
  complex_process2(classmate) # 複雑な処理2
  complex_process3(classmate) # 複雑な処理3
  score = classmate.get_score() # 例えば、成績を返したり
  all_score += score

# 全員の総合成績を出力
print(all_score)

とても楽ができました。

参考

Python のイテレータってなに？ | 民主主義に乾杯