【Python】pickleでリファクタリング・パフォーマンスチューニングのテストを簡単に書く

 通常、リファクタリングやパフォーマンスチューニングをする際は、特定のテストコードを準備して、コードを変更する都度テストやプロファイリングを行います。「パフォーマンスを計測するためのコード」も重要ですが、より重要なのは、「その変更によってインターフェースや動きが変わらないこと」を確認することです。

 既存のテストコードがあればそれを利用しても構いませんが、そもそもテストコードが十分でないことも多いですし、テストコードで現れないパフォーマンス上の問題があったりして、実際の動作に合わせたパフォーマンステスト用のテストコードを用意するのは非常に手間です。そんなときに使えるちょっとした技をご紹介します。

引数と結果のダンプを取る

 最も手っ取り早い方法は、引数と結果のダンプを取ることです。以下のコードを仕込み、モデルの状態などをファイルに書き出します。

import pickle

def add(a, b):
    with open('arguments.pickle', mode='wb') as f:
        pickle.dump((a, b), f)
    result = a + b
    with open('expected.pickle', mode='wb') as f:
        pickle.dump(result, f)
    return result

 これでargumentsとexpectedのファイルが書き出されます。これをテストコード上で読み取って、テストを回せばよいということになります。関数の最後にargumentsとexpectedを同時に取ることもできますが、Mutableなオブジェクトの場合、関数中で値が書き換わる可能性があるので、DeepCopyを行うか、今回のように分けて取ることをお勧めします。

def test_add():
    with open('arguments.pickle', mode='rb') as f:
        arguments = pickle.load(f)
    with open('expected.pickle', mode='rb') as f:
        expected = pickle.load(f)
    assert add(*arguments) == expected 

 テストが仕様に則っているかはともかくとして、ひとまず、コードの変更前とコードの変更後の整合性が取れるテストコードができました。これで中身を自由に弄れますね。

 テストケースが複数必要な場合は、以下のように連番をつけるのが得策です。

import pickle

i = 0

def add(a, b):
    global i
    i += 1
    with open('arguments' + str(i) + '.pickle', mode='wb') as f:
        pickle.dump((a, b), f)
    result = a + b
    with open('expected' + str(i) + '.pickle', mode='wb') as f:
        pickle.dump(result, f)
    return result

 これを仕込んでおき、テスト環境で結合試験を回しておけば、いつの間にかユニットテスト用のテストデータが出来上がっているはずです。インスタンスをそのまま塩漬けできるので、JSONにシリアライズしたりするより、手間がかからずに使い回すことができます。

まとめ

 pickleは、機械学習のモデルを保存するのによく使われたりしますが、意外なところで役に立ったりします。