doctest 駆動 AtCoder

AtCoder を始めた。ためしに Haskell や Nim で解いたりしていたが、結局は Python3 に落ち着いた。

解くのはいいが、サンプルのテストが面倒でしょうがない。TDD でできるといいが、フレームワークほどの仕組みは必要なさそう。その点で doctest はよかった。プログラムに埋め込めて、テスト実行も $ python -mdoctest XXX.py で済む。エイリアスを alias t='python -mdoctest' としておけば t XXX.py だけでテストが走る。

テンプレートはいろいろ調整して以下のようになった。 as_input の仕組みは他の Pythonista の回答を参考にさせていただいた。main 関数の引数がいびつだが、Fail 時にどのケースで失敗したかがわかるのでこうしている。このまま提出しても問題ないようになっている。

#!/usr/bin/env python

def main(_=0):   # 引数はとくに使わない。`as_input` を受け入れるためだけ
    N = int(input())
    print(N * (N + 1) // 2)

# string でテストケースを与えると stdin のように振る舞ってくれる
def as_input(s: str) -> None:
    import io
    global input
    f = io.StringIO(s)
    input = lambda: f.readline().rstrip()
    return None

sample1 = """5
"""
sample2 = """100
"""
def test():
    """
    >>> main(as_input(sample1))
    15
    >>> main(as_input(sample2))
    5050
    """
    pass

if __name__ == '__main__':
    main()

結果は (doctest -v) でこんなかんじにでてくれる。

$ alias t='python -mdoctest'
$ t -v test.py
Trying:
    main(as_input(sample1))
Expecting:
    15
ok
Trying:
    main(as_input(sample2))
Expecting:
    5050
ok
3 items had no tests:
    test
    test.as_input
    test.main
1 items passed all tests:
   2 tests in test.test
2 tests in 6 items.
2 passed and 0 failed.
Test passed.

解が浮動小数点の場合、とりあえずサンプルは実行するが、どうせ間違っているので目視でなんとなくあっていれば出している(適当)