Allow custom type checking via plugin for function definitions

まず get_function_hook と get_method_hook 周辺のプラグインインターフェースを読み、次に、呼び出しが解析される場所ではなく、関数定義の型チェックが行われる場所を追跡します。プラグインが各関数定義でカスタムの戻り値型エラーを報告でき、許可されたメソッドについては例外を引き続き可能にできれば、変更は完了です。issue では特定のテストファイルは指定されていません。

Feature

The possibility to allow custom type checking for function definitions using plugins would be very helpful to e.g. restrict allowed return types of a function. At the moment the plugin interface doesn't seem to have hooks for function definition type checking.

Pitch

Example: I want to restrict the return types of functions to allow them only to return a specific type.
Using get_function_hook/get_method_hook in a custom plugin like this

# -*- coding: utf-8 -*-
import mypy.errorcodes
import mypy.plugin
import mypy.plugins.default
import mypy.types

import typing

RESTRICT_UNIONS: typing.Final = mypy.errorcodes.ErrorCode(
    "restrict-unions",
    "some description",
    "custom_stuff")


class RestrictUnionsPlugin(mypy.plugin.Plugin):
    def get_function_hook(self, fullname: str) -> typing.Callable[[mypy.plugin.FunctionContext], mypy.types.Type] | None:
        return self.restrict_union_size

    def get_method_hook(self, fullname: str) -> typing.Callable[[mypy.plugin.MethodContext], mypy.types.Type] | None:
        return self.restrict_union_size

    @staticmethod
    def _check_union_size(ctx: mypy.plugin.FunctionContext | mypy.plugin.MethodContext, return_type) -> None:
        if isinstance(return_type, mypy.types.UnionType):
            if (type_count := len(ctx.default_return_type.items)) > 1:
                ctx.api.fail(
                    f"Too many union types: {type_count} > 1, {[str(typ) for typ in ctx.default_return_type.items]}",
                    ctx.context, code=RESTRICT_UNIONS)

    def restrict_union_size(self, ctx: mypy.plugin.FunctionContext | mypy.plugin.MethodContext) -> mypy.types.Type:
        self._check_union_size(ctx, ctx.default_return_type)
        return ctx.default_return_type


def plugin(version: str):
    # ignore version argument if the plugin works with all mypy versions.
    return RestrictUnionsPlugin

has the disadvantage that issues are raised on function call(s), not on definition.

E.g. running mypy with this custom plugin on the following piece of code

def some_method_with_union_return(some_input: str | int) -> str | int:
    return some_input


def some_other_method() -> None:
    some_method_with_union_return(5)


def another_method() -> None:
    some_method_with_union_return("foo")

finds 2 errors

source\main.py:6: error: Too many union types: 2 > 1, ['builtins.str', 'builtins.int']  [restrict-unions]
source\main.py:10: error: Too many union types: 2 > 1, ['builtins.str', 'builtins.int']  [restrict-unions]
Found 2 errors in 1 file (checked 2 source files)

This also means in case multiple return values are ok for certain methods, each place the method with multiple return types is called needs a # type: ignore[restrict-unions]