🐍

mypy pluginでオレオレルールを作ろう

2024/12/03に公開

Python

Advent calendar

Mypy

tech

こんにちは、こんばんは、おはようございます。LAPRASでエンジニアのような何かをしているdenzow です。この記事はLAPRAS Advent Calendar 2024 2日目の記事です。

mypyによる型チェック

Pythonで開発している環境なら最近はたいていmypyで型チェックを入れているのではないでしょうか。これがあるだけで本番にデプロイしてから思っても見なかった型エラーで死ぬことを防いでくれますが、標準のチェックルールに加えて開発チーム独自のルールを組み込んでいく方法を簡単に紹介します。

mypy plugin

mypyにはplugin機構があります。mypy.plugin.Plugin を継承し、独自のルールや解析結果の上書きなどを行うことができます。

Extending and integrating mypy

基本的には色々な解析のタイミングにHookを仕込めるようになっていて、そこで独自の処理を割り込ませるようにできています。

Hookの一覧

今回はget_class_decorator_hookを使って@dataclass デコレータをもったクラスでuser_codeという変数名がstrであることを強制するプラグインを作ってみます。

なお、完成しているサンプルリポジトリはこちらです。

Pluginのお作法

先程も少し述べましたが、プラグインはmypy.plugin.Pluginを継承して作ります。

sample_plugin.py

from mypy.plugin import Plugin

class UserCodeMustBeStrPlugin(Plugin):
    pass

def plugin(version: str):
    return UserCodeMustBeStrPlugin

スケルトンとしてはこんな感じで、実際のカスタムプラグインクラスとそれを戻すversion引数を受け取るplugin関数を定義します。

mypy実行時にこのpluginを組み込む場合はplugins に設定します。mypy.iniやpyproject.tomlで設定できます。

pyproject.toml

[tool.mypy]
python_version = "3.12"
plugins = "sample_plugin.py"

今回の解析対象のファイルです。

src/target.py

from dataclasses import dataclass


@dataclass
class ValidDataClass:
    user_code: str
    name: str


@dataclass
class InvalidDataClass1:
    user_code: int
    name: str


@dataclass
class InvalidDataClass2:
    custom_user_code: int
    name: str


class ExcludeMyPyTarget:
    user_code: int
    name: str

InvalidDataClass1/InvalidDataClass2 はstrにしたいuser_codeがintになっているのでこれを検知させることがゴールです。もちろん定義としては問題ないためプラグインなしではエラーは発生しません。

(mypy-sample) denzow@denzowmouse:~/work/denzow/mypy-plugin-sample$ poetry run mypy src/
Success: no issues found in 2 source files

get_class_decorator_hookの実装

get_class_decorator_hook を実装していきます。

sample_plugin.py

class UserCodeMustBeStrPlugin(Plugin):
    def get_class_decorator_hook(self, fullname: str):
        # decorator
        def decorator_hook(ctx: ClassDefContext):
            class_def = ctx.cls
            if fullname == 'dataclasses.dataclass':
                print(class_def.name)
        return decorator_hook

decorator_hook という ClassDefContext を受取り None を戻すHookを戻させています。これがmypyがデコレータを発見するたびに呼ばれます。(fullnameにはデコレータ名が入ってきます)。実行すると@dataclassのついたクラス名だけが表示されます。(ExcludeMyPyTargetが出てきていない)

(mypy-sample) denzow@denzowmouse:~/work/denzow/mypy-plugin-sample$ poetry run mypy src/
ValidDataClass
InvalidDataClass1
InvalidDataClass2
Success: no issues found in 2 source files

次に、この各クラスのattributesを取得します。このあたりは少し手探りですがこのような形で、属性名と型をlist[tuple[str,str]]として取得できます。

    def get_dataclass_attributes(self, class_def: ClassDef) -> List[Tuple[str, str]]:
        # dataclassの属性を取得する
        attributes = []
        for name, symbol in class_def.info.names.items():
            # symbol.nodeがVarであれば、それは属性を示す
            if isinstance(symbol.node, Var):
                var_node = symbol.node
                attr_type = var_node.type
                type_str = str(attr_type) if attr_type else "Unknown"
                attributes.append((name, type_str))
        return attributes

実際に表示させるとこのようになります。

ValidDataClass
[('user_code', 'builtins.str'), ('name', 'builtins.str')]
InvalidDataClass1
[('user_code', 'builtins.int'), ('name', 'builtins.str')]
InvalidDataClass2
[('custom_user_code', 'builtins.int'), ('name', 'builtins.str')]

mypy apiの呼び出し

ここまでくれば型チェックは難しくありません。user_codeを含む変数名の型がbuiltins.str であるかを見るだけです。ClassDefContextはapiというSemanticAnalyzerPluginInterfaceの属性を持っています。このapi経由で意味解析の機能を呼び出せるようになっていますが、この中にfailというメソッドがあります。これを呼び出すことでmypyのエラーを任意に発生させることができます。

sampple_plugin.py

class UserCodeMustBeStrPlugin(Plugin):
    def get_class_decorator_hook(self, fullname: str):
        # 全クラスのデコレータを調べるフックを設定
        def decorator_hook(ctx: ClassDefContext):
            class_def = ctx.cls
            if fullname == 'dataclasses.dataclass':
                attributes = self.get_dataclass_attributes(class_def)
                for attr_name, typer_str in attributes:
                    if attr_name.endswith('user_code') and typer_str != 'builtins.str':
                        ctx.api.fail(
                            f"Attribute '{class_def.name}.{attr_name}' should be of type 'builtins.str'",
                            ctx.cls
                        )
            return None
        return decorator_hook

attr_nameがuser_codeで終わる場合に、strでなければctx.api.failを呼び出してmypy上のエラーとしています。

実際の結果がこちらです。

(mypy-sample) denzow@denzowmouse:~/work/denzow/mypy-plugin-sample$ poetry run mypy src/
src/target.py:11: error: Attribute 'InvalidDataClass1.user_code' should be of type 'builtins.str'  [misc]
src/target.py:17: error: Attribute 'InvalidDataClass2.custom_user_code' should be of type 'builtins.str'  [misc]
Found 2 errors in 1 file (checked 2 source files)

期待通りにInvalidDataClass1/InvalidDataClass2をエラーとして処理できています。ちなみに、codeを指定していないのでmiscなエラーとして処理されていますが以下のように変更することで表示のカスタマイズもできます。

 # plugin.py
 from mypy.plugin import Plugin, ClassDefContext
 from mypy.nodes import ClassDef, Var
+from mypy.errorcodes import ErrorCode
 from typing import List, Tuple
 
 
@@ -16,6 +17,11 @@ class UserCodeMustBeStrPlugin(Plugin):
                         ctx.api.fail(
                             f"Attribute '{class_def.name}.{attr_name}' should be of type 'builtins.str'",
                             ctx.cls,
+                            code=ErrorCode(
+                                code='UserCodeMustBeStr',
+                                description='description',
+                                category='custom'
+                            )
                         )
             return None
         return decorator_hook

(mypy-sample) denzow@denzowmouse:~/work/denzow/mypy-plugin-sample$ poetry run mypy --show-traceback src/
src/target.py:11: error: Attribute 'InvalidDataClass1.user_code' should be of type 'builtins.str'  [UserCodeMustBeStr]
src/target.py:17: error: Attribute 'InvalidDataClass2.custom_user_code' should be of type 'builtins.str'  [UserCodeMustBeStr]
Found 2 errors in 1 file (checked 2 source files)

まとめ

いかがでしたでしょうか。mypy pluginを使うことでチーム独自の制約を定義、CIで矯正できることはPRレビューなどの効率化にも寄与できる可能性が有ります。正直ドキュメントがそれほど整っていないmypy pluginではありますが、小さな部分から検討していくのもいいかもしれません。

mypyによる型チェック

mypy plugin

Pluginのお作法

get_class_decorator_hookの実装

mypy apiの呼び出し

まとめ

Discussion