Allow custom type checking via plugin for function definitions

[waterfallwhitebread]

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]