itertools.chain() で連結

iterable なオブジェクトをまとめる itertools.chain() について紹介します。

itertools.chain() を使うことで、複数の iterableオブジェクトを連結することができ、一つの for文で複数のオブジェクトを扱えます。
以下は、文字列とリストを一つの for文 でまとめて走査する例です。

>>> from itertools import chain

# chain()の引数は文字列とリスト
>>> for i in chain('ABC', [1, 2, 3]):
...     print(i, end=(' '))
...
A B C 1 2 3

以下はもう少し実践的な、２つのフォルダの glob の結果をまとめて表示するプログラムです。

>>> from itertools import chain
>>> from glob import glob

# 複数のパス
>>> path1 = './doc/*.txt'
>>> path2 = './tool/*.txt'

# glob(path1)とglob(path2)をchain()で連結
>>> for i in chain(glob(path1), glob(path2)):
...     print(i)
...
./doc/readme.txt
./doc/changelog.txt
./tool/readme.txt

itertools.chain() には弟にあたる itertools.chain.from_iterable() もあります。
itertools.chain() は複数の引数を連結するものでしたが、chain.from_iterable()は引数を１つのみ取ります。
以下の chain() と chain.from_iterable() は同じ内容になります。

>>> chain('ABC', [1, 2, 3])

>>> chain.from_iterable(['ABC', [1, 2, 3]])

itertools.chain.from_iterable() の方が便利な場面があります。
以下は、先ほどの glob を使う例を改良したものです。

>>> from itertools import chain
>>> from glob import glob

# 複数のパスをリストで持つようにする
>>> paths = ['./doc/*.txt', './tool/*.txt']

# chain.from_iterable()の方がスッキリ書ける
>>> for i in chain.from_iterable(glob(path) for path in paths):
...     print(i)
...
./doc/readme.txt
./doc/changelog.txt
./tool/readme.txt