Pythonで型ヒントがついた関数のモックを作る

「テストのコードも例外なく型ヒントをつけたいなあ」

そんな気持ちでPythonを書いていたところ、型ヒントがついている関数のモックを作れないかなという興味からやってみたところなんとなくそれっぽいものができたのでご紹介します。

やりたいことのイメージとしては

function_mock = FunctionMock(引数).returns(返り値)
ret = function_mock(値) # ここでassertionが走る
use(ret) # 指定した返り値が帰っている

モック本体のコード

from __future__ import annotations

from typing import Callable, Generic, List, Optional, ParamSpec, TypeVar


Params = ParamSpec('Params')
Return = TypeVar('Return')


class Call(Generic[Params, Return]):
    retval: Optional[Return]

    def __init__(self, *args: Params.args, **kwargs: Params.kwargs) -> None:
        self.args = args
        self.kwargs = kwargs
        self.retval = None

    def returns(self, retval: Return) -> Call:
        self.retval = retval
        return self


class FunctionMock(Callable[Params, Return]):
    calls: List[Call[Params, Return]]
    count: int

    def __init__(self, calls: List[Call[Params, Return]]) -> None:
        self.calls = calls
        self.count = 0

    def __call__(self, *args: Params.args, **kwargs: Params.kwargs) -> Return:
        assert self.count < len(self.calls), f'Too many calls to function: want {len(self.calls)}.'
        call = self.calls[self.count]
        assert args == call.args, f'Argument mismatch: want {call.args}, got {args}'
        assert kwargs == call.kwargs, f'Keyword argument mismatch: want {call.kwargs}, got {kwargs}'
        self.count += 1
        return call.retval

    def close(self):
        assert self.count == len(self.calls), f'Not enough calls to function: want {len(self.calls)}, got {self.count}.'

    def __enter__(self):
        return self

    def __exit__(self, ex_type, ex_value, traceback):
        self.close()
        return ex_type is None

解説

関数のモックに型ヒントをつける際のポイントとして「引数の型ヒントをどうするか」が意外と曲者です。というのもPythonの関数はキーワード引数の存在によりただ単純に引数列をタプルとして比較できるわけではないのです。そこでParamSpecを使っていきます。

ParamSpec は関数の引数を表現するための特別な型パラメータを作るクラスで P = ParamSpec('P') という型パラメータを定義すると Callable[P, R] と関数の型を定義できたり P.args と P.kwargs という型で引数やキーワード引数に型ヒントをつけることができます。ただし一つだけ制約があって、P.args と P.kwargs はそれぞれ *args と **kwargs の型ヒントにしか使用できません。つまり事実上は関数の引数の型ヒントにしか利用することができません。そのためこれらの型ヒントを利用したい場合にはうまく関数を使って表現してあげる必要があります。この点に注意しつつ書いてきます。

型パラメータの定義

モックを作る上で引数と返り値の型を明示したいので引数の型は ParamSpec、返り値の型は TypeVar で表現します。

Params = ParamSpec('Params')
Return = TypeVar('Return')

関数呼び出しのオブジェクト化

関数のモックに求められる挙動として「呼び出された時の引数を突き合わせる」ことと「隠蔽した振る舞いに変わって値を返す」ことの2つがあります。先述の通り、ParamSpec の args と kwargs は関数の引数の型ヒントにしか使用できないためクラスのコンストラクタで使います。本当はコンストラクタで返り値も指定できるようにしたいのですが、ここで Callable[Params, Return] と違う型付けになってしまうのが微妙に感じたので別途 returns メソッドを追加してあります（いい回避方法があったらぜひご教示願いたいです）。

class Call(Generic[Params, Return]):
    retval: Optional[Return]

    def __init__(self, *args: Params.args, **kwargs: Params.kwargs) -> None:
        self.args = args
        self.kwargs = kwargs
        self.retval = None

    def returns(self, retval: Return) -> Call:
        self.retval = retval
        return self

モック本体

モック本体の挙動としては3つの挙動が期待されます。「呼び出された際に突き合わせる引数と返す値のリストを与えられる」こと、「呼び出されるたびに期待された引数が渡されていることを確認して指定の値を返す」こと、そして「指定回数だけ呼び出されたかどうかを確認する」こと。FunctionMock はこれらを愚直に実装しています。

class FunctionMock(Callable[Params, Return]):
    calls: List[Call[Params, Return]]
    count: int

    def __init__(self, calls: List[Call[Params, Return]]) -> None:
        self.calls = calls
        self.count = 0

    def __call__(self, *args: Params.args, **kwargs: Params.kwargs) -> Return:
        assert self.count < len(self.calls), f'Too many calls to function: want {len(self.calls)}.'
        call = self.calls[self.count]
        assert args == call.args, f'Argument mismatch: want {call.args}, got {args}'
        assert kwargs == call.kwargs, f'Keyword argument mismatch: want {call.kwargs}, got {kwargs}'
        self.count += 1
        return call.retval

    def close(self):
        assert self.count == len(self.calls), f'Not enough calls to function: want {len(self.calls)}, got {self.count}.'

    def __enter__(self):
        return self

    def __exit__(self, ex_type, ex_value, traceback):
        self.close()
        return ex_type is None

モックのテスト

使い方のデモを兼ねてpytestでテストを書いてモックが期待通りに動作しているかを確認します。

def test_function_mock():
    calls = [Call[[int], int](i).returns(i * 2) for i in range(10)]
    with FunctionMock[[int], int](calls) as double_mock:
        for i in range(10):
            double_mock(i)

まとめ

型ヒントで守られた関数のモックは作れる！その気になれば unittest.Mock みたいな高機能なモックも作れる気がしますが、ミニマリストとしてはこのくらいでいいんじゃないかという気がしています。パッケージ化するつもりは（執筆時点では）ないですが、反応次第では小さいパッケージとして公開するかもしません。ハッピー型ヒントライフ。