エムスリーテックブログ

エムスリー(m3)のエンジニア・開発メンバーによる技術ブログです

mypy plugin に入門して社内 OSS である gokart を型安全にしてみた

今回は mypy plugin を利用して、型安全に対応していないライブラリを型安全にする方法を紹介します！ 具体的にはエムスリーが開発する機械学習パイプラインツールである gokart を対象とし、mypy plugin を用いてどのように型の課題を解消したかについて解説します。 対象読者としては、既に gokart を使ってくださっている方はもちろんですが、dataclass や Pydantic がどのように型を担保しているかについて興味がある方も想定しています。

github.com

はじめまして！ 2024/07 から AI・機械学習チームにジョインした山本（@hiro-o918）です。 関西メンバーとして京都オフィスで働いています！ この記事はAI・機械学習チームブログリレー 5 日目になります。

入社して 1 ヶ月が経ち、徐々に会社やチームの雰囲気にも慣れてきました。 今回のブログリレーでは「入社エントリを...」とも思ったのですが、mypy plugin の作成に挑戦してみたのでその話をしたいと思います！

gokart について

gokart は、エムスリーが開発する機械学習パイプラインツールです。 次の記事に詳しく gokart について書かれているので、興味がある方はぜひご覧ください。

www.m3tech.blog

記事にも書かれている通り、gokart では SOLID 原則における Single Responsibility Principle を重視しており、パイプラインの中の各処理をタスクとして定義することで、処理の再利用性やテスト容易性を高めています。 今回はそのタスクに対する型安全性についての課題に取り組みました。

gokart における型の問題

クラス変数をハックするツールである

gokart では次のように記述することで、パイプライン上のタスク（一つの処理のまとまり）を定義できます。 近年 dataclass や Pydantic を利用したコードが主流になってきていることもあり、次のクラス定義とインスタンス化は非常に自然に感じます。

class MyTask(gokart.TaskOnKart):
    name: str = luigi.Parameter()

MyTask(name="mytask")

しかしながらPython 本来の文法を振り返って考えてみると、かなりの離れ業であることがわかります。 というのも、本来 Python においてクラス直下の変数はクラス変数でありインスタンス変数ではないためです。

例えば、次のようなコードを実行した場合、name をインスタンス変数として割り当てるというようなコンストラクタが存在しないため、エラーになってしまいます。

class Foo:
    name: str # ここの name はクラス変数

Foo(name="foo")

gokart を含めて先述のツールは継承元のクラスやデコレータの機能によって、本来クラス変数として定義されている変数をインスタンス変数として扱うことができています。

mypy による型チェックする上での課題

実行時においては各ツールがインスタンス変数までの割り当てをいい感じにやってくれることがわかりました。 一方で、このコードを静的に解析して型チェックをすることを考えたらどうでしょう？

理想的には mypy は次のコードを静的解析の結果、コメントに示しているように見分けないといけません。 実行時チェックではなく静的解析というのがキモで、このほとんど似たような変数定義を読み分ける必要があります。

class Foo:
    name: str 

# 異常なコード
# クラス変数でありコンスタラクタがないのでエラー
Foo(name="foo") 

@dataclass
class Bar:
    name: str

# 正常なコード
# dataclass において `name` はインスタンス変数であり、コンスタラクタも存在する
Bar(name="bar")

class MyTask(gokart.TaskOnKart):
    name: str = luigi.Parameter()

# 正常なコード
# gokart において `name` はインスタンス変数であり、コンスタラクタも存在する
MyTask(name="mytask")

これまで gokart では mypy がこの違いを見分けることができなかったため、コンストラクタに対する型チェックが行われていませんでした。 言い換えると、次のようなコードを書いても実行するまで気づくことができないという課題がありました。

MyTask(name=1)  # str に int を渡している

前職でも gokart を採用した経験があり、その当時から型安全性に対する課題を gokart に感じていました。 入社後社内でも gokart の型安全の話で盛り上がったので、その足がかりとして mypy plugin やってみるか！となったのが今回の機能開発のきっかけです*1。

mypy plugin を自作してみる

gokart における mypy の型チェック上の問題を示したところで、ここからは mypy plugin の概念に触れ、その仕組や開発方法について示していきたいと思います。 mypy には Plugin クラスがあり、これを継承して利用することでカスタマイズした処理を作成できます。

mypy plugin のコンセプトは静的解析時において、構文解析の結果を上書きすることにあります。 gokart のタスクにおける今回の問題は、クラス変数を利用したコンストラクタがコード上に存在しないことにあります。 よって今回の目標は、構文解析の結果を改変してコンストラクタを挿入すれば解決します。

class MyTask(gokart.TaskOnKart):
    name: str = luigi.Parameter()
    # この部分が構文解析の結果に追加されているとよい
    def __init__(self, name: str) -> None: ...

ここでは大きく 2 つのパートに分け継承後の Plugin クラスの実装から利用までの流れを示します。 最終的に、上記で示したコンストラクタを構文解析の結果に追加することを目指します。

  • 静的解析時のイベントにフックする Plugin クラス
  • 構文解析の結果を上書きする mypy API

静的解析時のイベントにフックする Plugin クラス

継承元である Plugin クラスは静的解析時のイベントにフックして実行される関数が用意されています。 各関数が担っている内容については、次のページで紹介されているので確認してみてください。

mypy.readthedocs.io

今回は get_base_class_hook() を利用します。 これは、あるクラスが継承を読み込んだときにそのクラス定義を書き換えるための hook になります。 この hook はなにかを継承しているクラスすべてで呼び出されるため、今回は継承元が gokart.TaskOnKart のときのみクラス定義を上書きするようにします*2。

class TaskOnKartPlugin(Plugin):
    def get_base_class_hook(self, fullname: str) -> Callable[[ClassDefContext], None] | None:
        sym = self.lookup_fully_qualified(fullname)
        if sym and isinstance(sym.node, TypeInfo):
            # この部分で継承元が 'gokart.task.TaskOnKart' であることを確認している
            if any(base.fullname == 'gokart.task.TaskOnKart' for base in sym.node.mro):
                # コールバックを返り値とすることで、このクラスに対しての構文解析の結果の改変する処理が呼び出される
                return self._task_on_kart_class_maker_callback
        return None

    def _task_on_kart_class_maker_callback(self, ctx: ClassDefContext) -> None:
        transformer = TaskOnKartTransformer(ctx.cls, ctx.reason, ctx.api)
        transformer.transform()

上記が gokart の mypy plugin における具体的なコードの一部です。 コード内のコメントにもある通り、継承元のクラス名を比較して目的のものを持っていた場合において、コールバックを返すようにしています。 これによって目的のクラスに対してのみ、構文解析結果の改変することが可能です。

構文解析の結果を上書きする mypy API

次に構文解析の結果を改変する処理を見ていきます。

class TaskOnKartTransformer:
    """Implement the behavior of gokart.TaskOnKart."""
    ...
    def transform(self) -> bool:
        """Apply all the necessary transformations to the underlying gokart.TaskOnKart"""
        info = self._cls.info
        # 1. クラス変数として定義されている変数を集約している
        attributes = self.collect_attributes()

        if attributes is None:
            # Some definitions are not ready. We need another pass.
            return False
        for attr in attributes:
            if attr.type is None:
        
    # 2. クラス定義に実際のコンストラクタがない場合において、改変処理を実行する
        if ('__init__' not in info.names or info.names['__init__'].plugin_generated) and attributes:
            # 3. `1.` で集めた変数を関数の引数となる形式に変換する
            args = [attr.to_argument(info, of='__init__') for attr in attributes]
            # 4. `add_method_to_class` を利用して、対象クラスの構文解析の結果に関数を追加する
            #      ここではコンストラクタを作成しており、その時の引数は `3.` の結果を利用している
            add_method_to_class(self._api, self._cls, '__init__', args=args, return_type=NoneType())
    # 5. おまじない。多重継承をしたときに子から親の変数も取れるようにするため、集約した変数を保存しておいて子の解析のときに参照する。
        info.metadata[METADATA_TAG] = {
            'attributes': [attr.serialize() for attr in attributes],
        }

        return True

コード内に書いている内容の再掲になりますが、次の流れで構文解析の結果にコンストラクタを作成しています。

  1. クラス変数として定義されている変数を集約している。
  2. クラス定義に実際のコンストラクタがない場合において、改変処理を実行する。
  3. 1. で集めた変数を関数の引数となる形式に変換する。
  4. add_method_to_class を利用して、対象クラスの構文解析の結果に関数を追加する。ここではコンストラクタを作成しており、その時の引数は 3. の結果を利用している。
  5. おまじない。多重継承をしたときに子から親の変数も取れるようにするため、集約した変数を保存しておいて子の解析のときに参照する。

以上の処理によって、自作のプラグインが gokart.TaskOnKart を継承した MyTask を解析する際に、def __init__(self, name: str) -> None: ... を持つようになりました！

class MyTask(gokart.TaskOnKart):
    name: str = luigi.Parameter()
    # この部分が構文解析の結果に追加された！
    def __init__(self, name: str) -> None: ...

mypy plugin 開発の Tips

最後に mypy plugin を開発する際に私が実際に行った方法を紹介します。

既存コードを参考にする

言わずもがなという感じですが、既存コードを参考にするのが走り出しとしては最も良いです。 特に、mypy 公式が提供している dataclass 向けの Plugin クラス は品質的にも間違いないためとても参考になりました。 また、Plugin クラスのコールバックである def get_base_class_hook(self, fullname: str) を GitHub で検索しても他のライブラリの実装を探すことができました。

後の残りはひたすら print と pdb を利用してデバッグをするのみです。 テスト対象の Python スクリプトを作成して、自作の mypy plugin を有効化した上で mypy CLI を呼び出します。 --pdb オプションは、プラグイン内部でエラーが起きたときにその場所でデバッガを起動してくれます。 また、--tb オプションを有効化することで、mypy plugin 内部のエラーも見ることができます（このオプションを有効にしないと、中で起きたエラーが握り潰されてしまいます）。

mypy test.py --pdb --tb

テストコードを追加する

最後にテストコードについてです。 構文解析のロジックが分かりづらいため、mypy を実際に使ったテストコードの追加は必須だと思います。 今回私は次のような感じでテストコードを作成しました。

class TestMyMypyPlugin(unittest.TestCase):
    def test_plugin_no_issue(self):
        # 1. lint をかけるコードの str
        test_code = """
import luigi
import gokart


class MyTask(gokart.TaskOnKart):
    # NOTE: mypy shows attr-defined error for the following lines, so we need to ignore it.
    foo: int = luigi.IntParameter() # type: ignore
    bar: str = luigi.Parameter() # type: ignore
    baz: bool = gokart.ExplicitBoolParameter()


# TaskOnKart parameters:
#   - `complete_check_at_run`
MyTask(foo=1, bar='bar', baz=False, complete_check_at_run=False)
"""

        with tempfile.NamedTemporaryFile(suffix='.py') as test_file:
            # 2. `1.` のデータを一時ファイルに書き出し
            test_file.write(test_code.encode('utf-8'))
            test_file.flush()
           # 3. mypy の API を呼び出し、このとき cache などのオプションは切るようにする
            result = api.run(['--no-incremental', '--cache-dir=/dev/null', '--config-file', str(PYPROJECT_TOML), test_file.name])
           # 4. コンソールの出力結果の比較
            self.assertIn('Success: no issues found', result[0])

まとめ

今回は gokart を通じて mypy plugin の開発方法について紹介しました。 実際に次のバージョンから gokart の mypy plugin が利用できるので、ぜひ試してみてください！

github.com

github.com

We are hiring !!

エムスリーでは社内外で実際に運用されている OSS があり、その設計から開発にチームで取り組むことができます！ 少しでも興味がある方は、次のURLからカジュアル面談をご応募ください！

jobs.m3.com

*1:dataclass や Pydantic が型安全にできているのだから、できるだろうというモチベーションでした

*2:厳密には luigi 側の処理も必要ですがここでは簡単のために省略します