競プロのためのPythonユニットテスト環境

システムプログラミングや低レイヤを触っていたり、いろんな自作XXを作る¹なかで「実装が面倒くさそうだと感じるハードルを下げたい。」と思うことが結構ある。言い換えると「大した問題ではない」と自分が思える範囲を広げたい。

そこでたまたま見かけた競技プログラミングの鉄則を買ってみた。この本を購入した一番の理由はPythonでのコード例があるという点。競プロ自体をやっていきたいわけではないので慣れてるPythonが思考しやすい。

で、実際に問題と解答例を見てみたら、競プロでは標準入出力を使う処理が一般的のよう。これは手元でサクサク試行錯誤するのが地味に面倒だなと思った。

ということで、次のような条件で使えるユニットテスト環境を用意した。

回答のコードはモジュールトップレベルにベタ書き
入力を標準入力から受け取る
回答を標準出力に書き出す
使うのは標準ライブラリだけ

ユニットテストの実行例

プロジェクトのディレクトリ構成はこんな感じ。A01.pyなどには回答のコードがモジュールトップレベルにベタ書きされているが、それ以外は一般的なPythonのパッケージ構成。

$ tree code
codes
├── __init__.py
├── A01.py
├── A02.py
└── tests
    ├── __init__.py
    ├── test_A01.py
    └── test_A02.py

実際にテストを動かすとこんな感じ。標準ライブラリのunittestを使ったごく普通のテスト。

# 全てのテストケースを実行
$ python3 -m unittest -v
test_answer (codes.tests.test_A01.TestCaseA01.test_answer) ... ok
test_answer (codes.tests.test_A02.TestCaseA02.test_answer) ... ok

----------------------------------------------------------------------
Ran 2 tests in 0.000s

OK

# 個別のテストケースを実行
$ python3 -m unittest -v codes.tests.test_A02.TestCaseA02
test_answer (codes.tests.test_A02.TestCaseA02.test_answer) ... ok

----------------------------------------------------------------------
Ran 1 test in 0.000s

OK

各テストケースの実装

一例として、テスト対象を次のcodes/A01.pyとする。

# codes/A01.py
N = int(input())
print(N * N)

この場合の具体的なテストケースcodes/tests/test_A01.pyは次のようになっている。

# codes/tests/test_A01.py
from codes.tests import BaseTestCase


class TestCaseA01(BaseTestCase):
    problem = "A01"
    params = [
        # 試したいパターンごとに要素を追加する
        # （入力テキスト, 期待する出力テキスト）
        ("5\n", "25\n"),
        ("10\n, "100\n"),
    ]

    def test_answer(self):
        super().base_test_answer()

クラス属性paramsが入力と出力の組みになっている。そのため、paramsの要素の数だけサブテストが回る仕組み。

テストケースの実装がこのようにシンプルなのは親クラスのBaseTestCaseで諸々の実装を行なっているため。

BaseTestCaseの実装

競プロ固有の面倒ポイントは「標準入出力を使う」と「モジュールトップレベルに書かれている」の2点。Pythonの仕組み上、モジュールトップレベルに書かれているコードはインポートするだけで実行されることになる。そのため、実装方針は次のようにした。

テストの実行時には標準入出力をリダイレクトする。
- 標準出力のリダイレクトは標準ライブラリcontextlib.redirect_stdout()として提供されている
- 標準入力のリダイレクトはredirect_stdout()を真似て実装する
  - 標準入力のリダイレクトは標準ライブラリでは提供されていない
  - 真似ると言っても具体的な実装は2行で済む
サブテスト実行ごとにテスト対象モジュールをインポートする
- テスト実行中は標準入出力がリダイレクトされている
- Pythonは実行中に同じモジュールを複数回インポートしないため、インポート済みならリロードする

from contextlib import _RedirectStream, contextmanager, redirect_stdout
from importlib import import_module, reload
from io import StringIO
from unittest import TestCase


class redirect_stdin(_RedirectStream):
    _stream = "stdin"


class BaseTestCase(TestCase):
    problem = ""
    params = []

    @contextmanager
    def redirect(self, input_text, expected_text):
        with redirect_stdin(StringIO()) as stdin, redirect_stdout(StringIO()) as stdout:
            stdin.write(input_text)
            stdin.seek(0)
            yield stdout
            stdout.seek(0)
            self.assertEqual(expected_text, stdout.read())

    def base_test_answer(self):
        mod = None
        for i, (input_text, expected_text) in enumerate(self.params):
            with self.subTest(i=i), self.redirect(input_text, expected_text):
                if not mod:
                    mod = import_module(f"codes.{self.problem}")
                else:
                    mod = reload(mod)

これで実際に動かしながら、サクサク試行錯誤できるようになった。

現状の実装では特定のパラメータのテストだけを試したいときにはクラス属性paramsをコメントアウトするか、別のテストケースに分ける作るしかない。とは言え、今のところはこれはあまり困っていないので放置してる。

オリジナルだと自作PEGパーサーrhoboro/toypegやゲームボーイエミュレータrhoboro/rustboyなど。自作OS、自作プログラミング言語がテーマの書籍の写経なども好き。 ↩