手触りの良いscript -- argparseとMakefileとiteratorの妙

手触りの良いscriptというものについて最近考えている。

ここでの手触りの良さとは何かというと粒度の自由さとアクセス性の良さに関するもの。別の言い方をするなら好きな形で扱え融通が効くようなもののこと。融通が効くというのはどういうことかと言えば、全体を一気に実行することもできるし、かと思えば一部分を小さく実行することもできるというようなもの。

この手触りの良いscriptということについて幾つか書きたいことがあったので連載のような形で記事を書こうと思う。

これは1つ目の記事。

あなたは何らかの入力を受取り、それを用いて何らかの指標を取得したい。ここで計算に要する時間はそう多くはなくつまるところ計算の高速化などについてはあまり求められていない。とは言え元の入力を利用して、それに加えて少しだけ何らかの値を混ぜ合わせるような付帯的な計算を行いたい。

例えば以下の様な計算(これは例なのですごく単純にしている。本来的には計算ではなく収集の部分に時間が掛かる)

x, y -> x', y', xy_rate

ここで、入力に対しての収集作業が以下の様に行われているとする(本当はこの時間がほとんど)

x' = x * x
y' = y

収集された結果からちょっとした付帯的な計算が行われる

xy_rate = x' / y

例えば以下のような入力から

x,y
49,-2
50,-1
51,0
52,1
53,2

以下の様な出力を返すことを考える。

id,x,y,x_y_rate
49,2401,-2,-1200.5
50,2500,-1,-2500.0
51,2601,0,-
52,2704,1,2704.0
53,2809,2,1404.5

繰り返しになるが、本来はもう少し時間の掛かる処理になる(例えば、x -> x * xなどではなく、URLからh1部分を抽出するであるとかそういう処理)。

収集作業

とりあえず、何かしら収集されるような作業がある。これをcollectという関数で表すことにする。イテレーターにしておくのが良い(それは後々触れる)。

def collect(itr):
    """何らかの収集の操作、元々はxとyを受取りそれを使って、何かのapiにアクセスしたりする"""
    for x, y in itr:
        # 本当はもう少し複雑
        x = int(x)
        y = int(y)
        yield {"id": x, "x": x * x, "y": y}

本来的には、入力として受け取ったx,yが例えばURLで、ヘッドレスのブラウザ的なものでリクエストするなどが収集作業ということになる。

付帯的な計算

付帯的な計算をちょっと加えたいことがある。そのような処理を収集作業と分けておくと見通しが良くなる。例えば比の計算などがそれに対応するものかもしれない。

全体が手軽に見通せる範囲なら収集でyieldされた結果をそのまま使えば良い。

def decoration_for_calculation(itr):
    for row in itr:
        # fetch()で取れる値がぱっと見で分かるようならyieldされた値に追加で計算した結果を付加して上げれば良い
        if row["y"] == 0:
            yield row
            continue
        row["x_y_rate"] = row["x"] / row["y"]
        yield row

名前についてはもう少し考えたほうが良いものの(長い)、言いたいことは関数名の通りで一種のデコレーターとして利用している。なので当然使うときには以下の様な形になる。

# 収集作業
itr = collect(itr)
# 付帯的な計算
itr = decoration_for_calculation(itr)

ところで収集作業の上で構成される値の構造が不明瞭な場合には、以下のように全てを記述して書くと見通しが良くなることもある。

def decoration_for_calculation(itr):
    for row in itr:
        # ここでどのような結果が欲しいのか全て網羅されるのが良い
        # つまり `id, x, y, xy_rate` という結果がほしい(おそらくCSV的な表現)
        r = {}
        r["id"] = row["id"]
        r["x"] = row["x"]
        r["y"] = row["y"]

        if row["y"] == 0:
            # ここで欠損値に値を追加するかは出力先の処理に依存する(csvのwriterはrestvalという値を指定すれば良いので追加せずyieldしている)
            yield r
            continue
        r["x_y_rate"] = row["x"] / row["y"]
        yield r

融通が効く呼び方

ここで呼び方について考えてみる。手触りの良さとは融通の効きやすさ、融通の効きやすさとは呼び方の多様さ。例えば以下のような多用な呼び方ができるようにする。

  • 収集処理が上手く動いているか調べたい
  • 収集処理中のインタラクションがまともかどうか調べたい
  • 収集後の計算の処理が上手く動いているか調べたい
  • 少なくとも1つのx,yのペアに対して上手く動くか調べたい

これらを試すために、コードを削ったり書き換えたりなどをしなければいけないのが融通が効かないコード。

テスト的な呼び方

例えば上の表現を別の言い方に揃えるとそれはテストの話になる。

収集処理が上手く動いているか調べたい

これは、収集部分のintegration test。

収集処理中のインタラクションがまともかどうか調べたい

これは、収集部分のunit test。例えばブラウザでのアクセスという話ならここをmockに置き換えての実行ということになる。

収集後の計算の処理が上手く動いているか調べたい

これは、付帯的な計算の部分のunit test。最もシンプルな形でTDDなどの練習ではこのあたりの要件で終わることが多い(?)。

少なくとも1つのx,yのペアに対して上手く動くか調べたい

これは、一種のE2E test。

そしてこれらについて、必要なら テストを書いておけるようにしたい。これは関数になっているのでわけない話(mockの部分をどうするかという話が不明瞭ではあるけれど、それはfetch()関数がfetcherのようなオブジェクトを取れば良いだけ。そこをmockに差し替える)。

ふつうの利用者としての呼び方

また、もう少し別の見方もできる。先程の表現はただの観点で、とりあえず実行してみることによって調べられれば十分ということもあるかもしれない。

つまり肝となるのはここの部分。

少なくとも1つのx,yのペアに対して上手く動くか調べたい

これはどういうことかと言えば、全てが繋がった処理を手軽に呼び出したいということになる。このために、所定のフォーマットでファイルを作って、所定の名前で保存してしてと言うようなことはしたくない(もっと言えば、所定のトリガーが発動されるまで待つなどもしたくない)。

例えば以下の様な形で呼べると嬉しい。

$ python fetch.py -x 50 -y -1
id,x,y,x_y_rate
50,2500,-1,-2500.0

この時入力のペアは1つだけになるので、お好みの場所でprintを差し込んでprintデバッグをしてみても良いし。あるいは好きな箇所でデバッガをあてて見たりしても良い。

少数の幾つかの入力に対しても試しに実行して見たい (new)

もう少し違った状況が考えられる。1つではないでもごく少数のペアを一度に処理して結果を見比べてみたいということもある。このようなときにも新たにファイルを作りたくない。全体ではなく部分を実行したい。

例えばコマンドがパラメータを複数受け取れるなら対応できる。

$ python fetch.py -x 49 -x 50 -x 51 -y -2 -y -1 -y 0
id,x,y,x_y_rate
49,2401,-2,-1200.5
50,2500,-1,-2500.0
51,2601,0,-

この時、実行結果をtraceできたりするととても便利。例えば、URLへのアクセスにproxyを加えて全てのrequest/responseが全部別途出力される。出力された中をgrepしてみたりして収集が上手く言っているか調べたい。

よりたくさんのペアを入力として一気に実行した結果が欲しい (new)

これは1つだけの実行の対極を考えることもできるかもしれない。1つのペアに限らずN個のペアを入力として一気に実行した結果が欲しいということ。というよりも実際の用途でこれこそが本懐と言っても良い。入力としてCSV(とは限らないが何らかのファイル)があり。それを実行した結果もまたCSV(とは限らないが何らかのフォーマットのファイル)で返す。というような。

$ python fetch.py --input input.csv
id,x,y,x_y_rate
49,2401,-2,-1200.5
50,2500,-1,-2500.0
51,2601,0,-
52,2704,1,2704.0
53,2809,2,1404.5

先頭N件だけ実行したい (new)

先程の少数のN個の類型として、先頭n件取り出すというようなこともできると良い。

$ python fetch.py --input input.csv --size=3
id,x,y,x_y_rate
49,2401,-2,-1200.5
50,2500,-1,-2500.0
51,2601,0,-

とりあえず動いていそうということを確認するのにはせいぜい5個くらいが動いていれば大丈夫なことも多い。

今度はもう少し大きな範囲で考えてみることにする。

一気に実行したいと言っても、それが全てではないかもしれない。例えば本番環境とステージング環境が分かれていて、手元でざざっと確認したい場合にはステージング環境のものを、実際には本番環境のものを入力として使いたいということもあるかもしれない。これらに対していちいち引数を変えてなど右往左往したくない。

何らかのflavorでタグ付けされた形で一気に実行したい (new)

これは実質的には手動で以下のように書けば良いということになるが。

$ python fetch.py --input production-input.csv | tee production-output.csv
# もしかしたら標準出力で返すのではなく以下の様な形のものもあるかもしれない
# $ python fetch.py --input production-input.csv --output production-output.csv

$ python fetch.py --input staging-input.csv | tee staging-output.csv

これらを手動で書きたくない。どうせなら何らかのタスクとして一連の動作を動かしたい(後述のタスクランナーをどうするかという話につながる)。

また、人間はミスをするし、開発者は人間なので、作った処理が全部まともに動くかどうかわからない。なので処理を再開したい。

失敗した所からresume(retry)したい (new)

つまるところ、以下の様なことができると良い。

$ python fetch.py --input input.csv
# ...
# 404 errorが発生
# コードを直して、再開する
$ RESUME=1 python fetch.py --input input.csv

これに関しては以前記事を書いた

このようなわがままな状況に良い感じに対応するのが手触りの良いscriptということになる(ちなみにこの記事では環境がローカルだけ保存先も手元のCSVということにしているけれど。これがクラウド上のバッチ処理でDBに保存だったり、FaaS上でpubsub的なイベントをトリガーに実行されたりという変形がある(今回は触れない))。

融通を効かせる

ある程度どのようなものが手触りの良いscriptかということが説明できたと思うので実装に移る。実際に融通の効くscriptを作りたい。

1個のペアの入力に対応する

1つのペアの入力に対応するのは、単なるコマンドライン処理argparseで雑に対応できる。

import argparse

parser = argparse.ArgumentParser()
parser.add_argument("-x")
parser.add_argument("-y")

# $ python fetch.py -x <x value> -y <y value> で args.x, args.y に値が入る
args = parser.parse_args()

N個のペアの入力に対応する

でもちょっと待って欲しい。上の例はあまり良くない。というのもN個のペアに対応する場合の条件分岐が出てくる。なので常にN個受取りその特殊な状態が1個という風に考えると良い。幸い標準ライブラリのargparseであってもそれには対応できる。

import argparse

parser = argparse.ArgumentParser()
parser.add_argument("-x", action="append", dest="x_list")
parser.add_argument("-y", action="append", dest="y_list")

# $ python fetch.py -x <x value> -y <y value> で args.x_list, args.y_list に値が入る
args = parser.parse_args()

ここでiteratorにしておいたことが生きてくる。実際の実行では以下のようにzipしてあげれば良い。

def run(*, config, x_list=None, y_list=None):
    r = zip((x_list or []), (y_list or []))
    itr = collect(r)
    itr = decoration_for_calculation(itr)
    yield from itr


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("-x", action="append", dest="x_list")
    parser.add_argument("-y", action="append", dest="y_list")

    itr = run(**vars(args))
    w = csv.DictWriter(sys.stdout, fieldnames=["id", "x", "y", "x_y_rate"], restval="-")
    w.writeheader()
    w.writerows(itr)

(ちょっとしたtips。argparseでparseした結果をvars()でdictにしてあげてrunを呼ぶようにすると、コマンドライン引数の受け渡しが明示的かつ手軽になって良い)

ファイル入力に対応する

コマンドライン引数だけで対応していたのではN個が100個などになってきた時に辛くなる。当然入力方法としてファイルなどからできる必要がある。

import argparse
parser = argparse.ArgumentParser()

# $ python fetch.py -x 10 -y 10 に対応
parser.add_argument("-x", action="append", dest="x_list")
parser.add_argument("-y", action="append", dest="y_list")

# $ python fetch.py --input input.csvに対応
parser.add_argument("--input", type=argparse.FileType("r"))

場合分けができれば良い。なのでrunは以下の様になる。

def run(*, config, x_list=None, y_list=None, input=None):
    if input is not None:
        r = ((row["x"], row["y"]) for row in csv.DictReader(input))
    else:
        r = zip((x_list or []), (y_list or []))

    itr = collect(r)
    itr = decoration_for_calculation(itr)
    yield from itr

先頭N個に対応する

先頭N個の対応はitertools.isliceが良い。ここでiteratorにしていたことが生きてくる。itertools.isliceはslice操作(e.g xs[x:y]などのこと)をiteratorに対して提供してくれるもの(resume対応などを考えるとiteratorを一括で全て消費してlistなどにしてしまうよりiteratorのままsliceしたほうが嬉しい)。

def run(*, config, size=None, x_list=None, y_list=None, input=None):
    if input is not None:
        r = ((row["x"], row["y"]) for row in csv.DictReader(input))
    else:
        r = zip((x_list or []), (y_list or []))

    if size is not None:
        r = itertools.islice(r, size)

    itr = collect(r)
    itr = decoration_for_calculation(itr)
    yield from itr

argparseに以下を追加。

parser.add_argument("--size", type=int)

融通を効かせつつも手軽さを

多用なコマンドラインオプションで融通を効かせつつも手軽さを求めたい。具体的には個別にコマンドに渡す引数を入力したくない。そういう場合にタスクランナーを使う。何を使っても良いが今回はおそらくどこにでもあるだろう消極的な理由でmakeを使うということにする(makeが最高とは全く思っていないけれど。polyglotな環境で試す時にある言語のデフォルトのタスクランナーを強制するというのがなかなか難しい。政治的な駆け引きなどしたくないし。あるツールのインストールを強制するというほどの恩恵を持つものが存在しない。環境によっては母国語が存在しそれへのデファクトがあるならそれを使えば良い(これはpythonに関しても同様の話))。

ちょっとした確認を

1個だけの実行、N個だけの実行を試して上手く動くか確認したい。雑にoneとmanyというタスクを用意する。

CMD := python fetch.py -c config.json

one:
  ${CMD} -x 1 -y 2

many:
  ${CMD} -x 1 -x 2 -x 3 -x 4 -x 5 -y 2 -y 1 -y 0 -y 1 -y 2

そしてコードを書き換える度に以下の様にして実行結果を確認する。

# 大丈夫そうかな?
$ make one
# たぶん大丈夫そう。もう少し多めの入力で試す
$ make many

一気に実行を

当然ではあるけれど、全部まとめて一気に実行したい。inputとなるcsvの名前を決めておけば手軽にまとめて実行できる。

CMD := python fetch.py -c config.json

default:
  mkdir -p dist
  ${CMD} --input ./input.csv | tee dist/output.csv

defaultタスクにしておけば以下で済む。

$ make

# ちなみにどのようなコマンドが実行されるのか不安なら -n 付きで呼び出す(実行されるであろうコマンドが表示されるだけ)
$ make -n
mkdir -p dist
python fetch.py -c config.json --input ./input.csv | tee dist/output.csv

幾つかのflavorを

ステージング環境に対する実行と本番環境に対する実行を呼び分けたい。これはMakefileのdefaultの変数をoptionalにして実行時に環境変数として設定することで対応できる。

PREFIX ?= test
CMD := python fetch.py -c config.json

default:
  mkdir -p dist
  ${CMD} --input ./${PREFIX}-input.csv | tee dist/${PREFIX}-output.csv

呼び分けられる。

# テスト環境
$ make
# 本番環境
$ PREFIX=production make
# ステージング環境
$ PREFIX=staging make

先頭N件

先頭N件だけという実行もタスクの引数にしてしまえると便利。ちょっとトリッキーではあるがmakeの関数を使うことで対応できる。

PREFIX ?= test
SIZE ?=
CMD := python fetch.py -c config.json $(addprefix --size=,${SIZE})

default:
  mkdir -p dist
  ${CMD} --input ./${PREFIX}-input.csv | tee dist/${PREFIX}-output.csv

SIZEを指定したときだけ --size ${SIZE} がオプションとして渡される。

(ちなみにそういう意味ではone,manyのタスク定義ででてきたmanyの部分もaddprefixを使って楽に記述ができる)

xs := 1 2 3 4 5
ys := -2 -1 0 1 2

many:
  ${CMD} $(addprefix -x ,$(xs)) $(addprefix -y ,$(ys))

最後に

全部つなぎ合わせたgistです。 眠れなかった結果、謎のそこそこ長大な文章が生まれた。