Python: Generator のレシピ

最近 python を使い始めたが generator、これは良い。不慣れな python ではあるが generator を作ったり使ったりした際に気付いた事などを纏めておこうと思う。 書いてみたらえらく長いエントリになってしまって、纏めたとは言い難いかもしれないが。

generator の前に iterator について

generator は iterator とよく似ていて、他の言語でも iterator は馴染みがあるので、iterator について他の言語との違いなどの印象を少し。 コレクションから iterator を取得する時、C++ で言う所の std::begin() としては iter() を使う様だ。尤も python でも大抵のコレクションは、in の後ろで評価されると iterator を返すし、iterator を引数とする関数にコレクションを渡せば iterator 扱いしてくれるようなので、ループの前に初期化して置く必要がある場合に使う位かなと思う。 第二引数(=sentinel)を渡す使い方は、終了条件を必要とするループでもシンプルに書けて便利。

python の iterator は次の要素が有効かどうかを判別するための関数、他の言語で言う bool 型の値を返す hasNext() の様なインスタンスメソッドは無く、最後迄繰り返した後は StopIteration 例外を raise してループを抜けるのが流儀だそうな。終端の判定代わりに例外を投げるというのは、余り私の好みでは無いので他のやり方を散々探してみたが、そういった方法は見つけられず、python ではそれが正義らしい。 そして、次の要素への iterate を行うには next(itr) を使う。iterator が next() メソッドを持つ訳ではない。いや、一応有るにはあるが、__next__ というのはライブラリの実装側が使うもので、普段使いするものではないと思っている。

あと、itertools というモジュールは iterator を扱う上で非常に便利な関数が提供されているので、どの様なものが用意されているか、一通り確認しておいたほうが良い。

python の generator

yield で値を返す関数を定義するか、generator 式を記述すれば定義できる。generator は遅延評価されるので、無限ループを定義して使っても、呼び出し側で呼び出しをキチンと止めれば問題にはならない。 なお、yield を使うとき、C# の癖で yield return と書くと文法エラーになる。

generator 式は配列内包表記の角括弧を丸括弧にするだけというお手軽さなので、そこらじゅうに書き散らかせる。だが、丸括弧は関数呼び出しとか、優先度を上げる演算子でもあるので、混在させるとコード読解の難易度は容易く上がっていく。generator は一旦その処理を表す為の適切な名前が付けられたローカル変数にでも代入しておいて、その変数経由で使用した方が、コードの読み手に優しいであろう。遅延評価されるのが基本の generator はこの様なエイリアスを作っても、その場で処理が動く訳ではないので、大したオーバーヘッドにはならないだろう。 それに python では長くなってしまった式の途中で改行しようとすると、行末にバックスラッシュエスケープを置いておく手間と、それによって少しコードがブサイクになる事を容認する広い心が必要になる。

文法

x for x in (GENERATOR) という文法について。 for の後ろは generator から返される値の一時変数名宣言、for の前でその一時変数を利用する式を記述する。C# の Linq 式で表せば from x in (GENERATOR) select x と同等の式になる。 tuple を返す generator だと構造化束縛（python での呼称は知らぬ）したりも出来る。

yield from

generator を使いはじめると、ある generator から取り出した値を評価して、その結果を yield で返す generator を作るといった事が多くなる。 この様に generator を連鎖的に利用する場合、yield from に置き換えるとコードがスッキリする。

for i in (f(x) for x in anIter):
  yield i
# ↑の様なものは、↓のように書き直せるという事
yield from (f(x) for x in anIter)

代入式

python 3.8 で導入された機能。pascal の様な代入演算子を使うと generator 式の中でもその変数を介して計算結果を参照出来るようになる。その演算子(:=)はセイウチ演算子とも呼ばれるらしい。なるほど似てるが私には寧ろバカボンのパパ ||:3= を想起させる。 3.8 より前だと、generator の for の前にある結果を if 条件式で使いたい場合には、generator をネストさせたり、適当な関数（filter とか）を使って if 条件式を除去したり、無理矢理感が漂うコードになったが、これを使えば解消出来る。 公式What's new に挙がっているサンプルは使い所がわかりやすい。

例題

以下、使用例を列挙。

例 1

ファイルを開いて一行ずつ読み込み、空白文字区切りの小数点数値を配列にして返す generator

def readVals(fname):
  with open(fname, "r") as instr:
    yield from ([float(v) for v in line.split()] for line in instr)

python ではテキスト入力ストリーム(instr)を iterate すると、一行単位の文字列(line)を要素としてループするようだ。ここでは空白区切りでテキスト分割するという仕様にしたので、引数なしの split メソッドで返される文字列(v)のリストを起点として、文字列から浮動小数点型に変換する様に評価する内包表記を書けば、後はそれを yield from で返せば良い。

例 2

与えられたファイル名(fname)のファイルをバイナリとして読み込み、指定サイズ(size)の byte 配列をその都度返す。

from functools import partial

def readChunk(size, fname):
  with open(fname, "rb") as instr:
    yield from iter(partial(instr.read, size), b'')

functools.partial() は関数に部分適用を行い、その部分適用された関数を返す関数。 read() はもともと呼び出す度にファイルポインタを進めて、前から順にイテレータのような動作をしてくれる為、部分適用しておいて引数を取らない関数として iter() に渡せば、 呼び出されるたびに指定サイズ分のバイト配列を前から順に返すイテレータとして扱えるので、yield でその byte 配列を返せば generator が定義できる。 終了条件としては、ファイル終端では空の byte 配列が返って来るので、iter() 関数の第二引数に空バイト配列を sentinel として指定しておけば、EOF に到達した時点でループが終了する。

例 3

文字列を一定文字数で分割する generator

from itertools import takewhile, repeat, islice

def splitBySize(size, text):
  src = iter(text)
  yield from takewhile(bool, ("".join(islice(src, size)) for _ in repeat(None)))

# test
src = "abcdefghijklmnopqrstuvwxyz" * 10
for t in splitBySize(10, src):
  print(t)

最初に入力(text) iterator を初期化する。generator 式の評価中で呼び出す（islice(iter(src), size) for...という風に）と、毎回シーケンスの最初を指す iterator に初期化されてしまい、無限ループを作ってしまう。 generator 式の初期化は in の後ろの部分だけと覚えておけば良さそう。

islice は iterator の現在位置から指定回数分繰り返して要素を返す iterator を返す。これを join で連結して文字列化すれば所望の要素が一つ完成する。だが islice を単発で呼び出してもシーケンスの最後まで処理される訳ではない。だから for 以降の無限ループで繰返し評価して値を取得する。 無限ループを使うと当然ながらシーケンスの終端を過ぎても空文字列を返し続けるので、takewhile の述語関数として bool() を指定して終了条件とする。

なお、string をストリーム扱いできるモジュール StringIO を使えば、例 2 と同じ様に書ける

from io import StringIO

def splitBySize(size, text):
  with StringIO(text) as instr:
    yield from iter(partial(instr.read, size), "")

この方がシンプルかと思うが、takewhile() を使う方法はループ終了の述語として、引数を一つ取る関数を渡せるので応用力は高いと思う。

例 4

バイナリシーケンスを読んで、byte 配列に分割して返す generator 関数が欲しい。但し分割する長さは一定ではなく、長さをその都度返す所与の generator から取得する事。 入力シーケンスが終端に達した場合、読込対象シーケンスの長さに合わせたバイトの長さに合わせる（固定長としてパディングしたくないという事）。

from itertools import takewhile, islice, cycle

src = "abcdefghijklmnopqrstuvwxyz" * 10
srcb = src.encode("utf-8")
sizes = (int(n) for n in cycle("123456789"))

def szsplitB(src: bytearray):
  itr = iter(src)
  yield from takewhile(bool, (bytes(islice(itr, i)) for i in sizes))
    
for s in szsplitB(srcb):
    print(s.decode())

例 3 str の時とそんなに変わらないが、python 初心者の私には byte を返す iterator から byte 配列を作る方法になかなか到達できず（結局 bytes() を呼びだす事になった）難儀した記念。ByteIO を使うのも良いと思う。 それ以外は例 3 の時に repeat() で無限ループにしていた部分を cycle() を使って整数を返し続けるようになっているだけ。

例 5

あるディレクトリ以下にあるファイル名を正規表現パターンで検索したい。

import re
import os

def matchedFiles(dir: str, pattern: re.Pattern):
  for rtdir, _, files in os.walk(dir):
    yield from (os.path.join(rtdir, f) for f in files if pattern.search(f))

# sample
for f in matchedFiles("..", re.compile(".+\.json")):
  print(f)

for の後ろに if で条件式を書けばフィルタとして使えるので、正規表現がマッチしたかどうかを返す式を書けばファイル名をフィルタすることができる。 generator とは関係ないが、python で pattern.match(f) とすると、文字列の先頭からマッチしていないと true が返らないというのはハマりポイントだった。

例 6

あるディレクトリ(path)のログファイル名にシーケンス番号が付いている、その中で一番大きなシーケンス番号を取得したい。

xxx-1.log
xxx-2.log
...
xxx-10.log
xxx-11.log

というようなファイルが有ったとしたら、11 を取得したい、という事。

import re
import os
from glob import glob
from os.path import join

def lookupMaxseq(path):
  nPat = re.compile(r"xxx-(\d+)\.log")
  fnames = (nPat.search(name) for name in glob(join(path, "xxx-*.log")))
  return max((int(m.group(1)) for m in fnames if m != None), default=-1)

ファイル名を正規表現でマッチさせて、番号部分の数値を取り出して、max 関数で評価させるという方針でやってみる。 正規表現パターン(nPat)は複数回使い回すので、compile を呼んでおく。パターン自体の説明は簡単なので割愛。 glob() を使えば、ファイル名リストへのイテレータが返ってくるので、それをパターンマッチさせた結果を要素とする generator を変数に入れておく(fnames) 。

この fnames を元にループさせて、マッチ結果オブジェクト(m) をみて、None だったらマッチしなかったのでスキップし、マッチしていたら数値部分とマッチさせた結果が group(1) に文字列として入っているので、それを整数に変換すれば所望の要素が作られる generator が書けた。

今度はその generator を max 関数に渡す。max は引数一つで呼び出した時は、要素が空の場合に例外を投げるので、default パラメータで空の場合に返される値を指定する。取り敢えず -1 でも返しておけば良いだろう。

例 7

文字列エンコード・デコードについて。細切れに読み込んだ byte 配列を utf8 文字列に変換する処理と、それとは逆に utf8 を byte 配列に変換する処理。

普通は読み込み時点で utf-8 扱いでテキスト入力するだけだが、データを byte 配列として読み込み、何かしらの変換を掛けてから、utf8 にする様な処理というのは時々必要になる。

mojiretu.encode("utf-8") とかで間に合う場面であればそれで十分なのだが、実務的には入力内容をまるごとメモリに載せていい場合というのは少ない（特に攻撃を受ける可能性が 0 ではない場合）ので、入力バッファへ（例 2 の様に）取り込んで逐次的に処理する場合が多い。 バッファの中身を単純に inbuf.decode("utf-8") して順繰りに文字列をつなげると、マルチバイト文字の途中でバッファが分割された場合に文字化けが生じてしまう。出力も同様の問題が生じる。 python ではそういう場面の為に iterdecode/iterencode が用意されている。 下のコードは使い方の例示のため、読み込んだバイト列をそのまま decode するが、バイト列に何も手を加える処理が無いのであれば、普通に open() で読めばいい。

from codecs import iterdecode, iterencode
from itertools import takewhile, islice, chain, repeat

def readUTF8byte(name):
  yield from iterdecode(readChunk(64, name), "utf-8")

# itr は不定長の byte 列を順次返す iterator。
# 入力を一旦連結して size で指定した大きさの byte 配列を順次返す。
def splitEvery(size, itr):
  flat = chain.from_iterable(itr)
  yield from takewhile(bool, (bytes(islice(flat, size)) for _ in repeat(None)))

# これはファイルを読んでるだけだが、実際にはなにか処理を行って
# 文字列(utf8)を順繰りに返す generator を想定する
def awesomeTaskResult():
  with open("sample.txt", encoding="utf-8") as istr:
    yield from istr

with open("sample_out.txt", "wb") as fout:
  for buf in splitEvery(64, iterencode(awesomeTaskResult(), "utf-8")):
    # 固定サイズ byte バッファに必要な処理があればここで行う。
    fout.write(buf)

readUTF8byte は、例2 の readChunk() を使って 64 バイトずつバイトストリームを読み込んで、そのバッファの内容を iterdecode() を使って utf8 文字列にして返す generator。それだけ。

iterencode() は逐次的にバイト列にして返してくれるので、それをそのまま出力ストリームに write() すれば良いのだが、 それだけだと普通に出力ストリームを encoding パラメタ付きで open() すればいいという話として終わってしまうので、 もう少しシナリオを加えて、出力の前に固定長のバイト配列を一旦作り、圧縮とか暗号化なんかの処理を行う必要があるとしよう。

なにか文字列として処理を行った結果の generator を iterencode() に渡す。この関数は iterator を返すので、splitEvery() 関数に渡す。 ここで使っている chain.from_iterable() は、iterator が返す細切れになった byte 列を一つにつなげて、1 byte ずつ順次返す処理をしてくれるので、例 3 と同様の generator を使えば、指定したサイズの byte 列（もちろん最後に返ってくるバッファのサイズはデータ末尾までに切り詰められる）が変えるので、それが for ループの buf 変数に代入する。 buf 変数を使ってなにか処理をして write() すれば、要求の出力が得られる。