yiskw note

機械学習やプログラミングについて気まぐれで書きます

【Python】MonkeyTypeを活用して型アノテーションを自動で追加する


MonkeyTypeとは

今回は、Pythonの型アノテーションを自動で付与するMonkeyTypeを使用してみました。
MonkeyTypeは、実行時の引数や返り値の情報から、自動でスタブファイルを生成したり、型アノテーションを追加してくれるライブラリです。
※ スタブファイルとは、型の情報を保持したファイルで、型検査時に用いられるもの。(参考: PEP 484 – Type Hints | peps.python.org)
※ 似たようなツールとして、pyannotateがありますが、こちらは長いことメンテナンスされていないようです。

MonkeyTypeは、型アノテーションを付与していない既存プロジェクトに、型アノテーションを導入する手助けになったり、 型アノテーションが付与されていないOSSのコード(深層学習の論文の実装など)を解読するのに役立つと思い、 使い方や使用感についてこちらにメモを残しておこうと思います。

実行環境

  • Python 3.9.9
  • MonkeyType 22.2.0

インストール

pipで簡単にインストールすることができます。

pip instsall MonkeyType

使い方

以下のようなmodule.pymain.pyの二つのファイルを用意します。

module.py

from numpy import ndarray
from typing import Optional, Union

def add(a, b):
    return a + b


def multiply(a, b):
    return a * b


def func_with_optional(a=None):
    if a is None:
        return 0

    return a ** 2

def create_dict_from_list(list_):
    dict_ = {i: val for i, val in enumerate(list_)}
    return dict_

main.py

import numpy as np

from module import add, func_with_optional, multiply, create_dict_from_list


def main():
    add(1, 2)

    a = np.ones((1, 2))
    b = np.ones((1, 2))

    multiply(a, b)
    multiply(1, 2)

    func_with_optional()
    func_with_optional(10)

    list_ = [i for i in range(10)]
    create_dict_from_list(list_)


if __name__ == "__main__":
    main()

以下を実行することで、実行時の型の情報をSQLiteデータベース(./monkeytype.sqlite3)に保存してくれます。

$ monkeytype run main.py

次に以下を実行することで、SQLiteのデータベースからスタブファイルを生成します。

$ monkeytype stub module
from numpy import ndarray
from typing import (
    Dict,
    List,
    Optional,
    Union,
)


def add(a: int, b: int) -> int: ...


def create_dict_from_list(list_: List[int]) -> Dict[int, int]: ...


def func_with_optional(a: Optional[int] = ...) -> int: ...


def multiply(a: Union[int, ndarray], b: Union[int, ndarray]) -> Union[int, ndarray]: ...

最後に以下を実行することで、スタブファイルを元に、型アノテーションを追加することができます。

$ poetry run monkeytype apply module

module.pyを見てみると、ちゃんと型アノテーションが追加されているのが確認できました。

from numpy import ndarray
from typing import Dict, List, Optional, Union

def add(a: int, b: int) -> int:
    return a + b


def multiply(a: Union[int, ndarray], b: Union[int, ndarray]) -> Union[int, ndarray]:
    return a * b


def func_with_optional(a: Optional[int]=None) -> int:
    if a is None:
        return 0

    return a ** 2

def create_dict_from_list(list_: List[int]) -> Dict[int, int]:
    dict_ = {i: val for i, val in enumerate(list_)}
    return dict_

使ってみての感想

上記結果を見て分かる通り、単一の型だけでなく、UnionOptionalにも対応しており、実行時の挙動に忠実に型アノテーションを付与してくれます。
また型アノテーションに必要なクラスは、自動でインポートしてくれる点も非常に便利です。
このように、コードを実行するだけで、手軽に型アノテーションを付与できるため、 既存プロジェクトに型アノテーションを導入したい場合や、型アノテーションのないコードを解読したい場合に役立ちそうです。

一方で気になった点として、Python3.9を使用していたとしても、リストや辞書のアノテーションにはtyping.Listtyping.Dictが使用されてしまいます。
特にIssueも上がっていなかったので、時間があれば自分の方でPRを出してみようかと思います。

まだまだ使いこなせていませんが、かなり便利なツールですので、今後も色々使用していきたいです。

参考