Allow custom type checking via plugin for function definitions

Bắt đầu bằng cách đọc các interface plugin xung quanh get_function_hook và get_method_hook, sau đó lần theo nơi các định nghĩa hàm được kiểm tra kiểu thay vì nơi các lệnh gọi được phân tích. Thay đổi hoàn tất khi một plugin có thể báo cáo các lỗi kiểu trả về tùy chỉnh tại mỗi định nghĩa hàm, trong khi các ngoại lệ đối với những method được cho phép vẫn có thể thực hiện; issue không nêu tên tệp test cụ thể nào.

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]