getattr_plugin.py

from mypy.plugin import Plugin, FunctionContext
from typing_extensions import Protocol
from typing import TypeVar, Any
from types import new_class


def handle_function_call(function_ctx: FunctionContext):
    T = TypeVar('T', bound=Any)
    attr = function_ctx.args[0][0].value

    def _update_annotations(ns: dict) -> dict:
        ns['__annotations__'] = {attr: T}
        return ns

    target = new_class("Target", (Protocol[T],), {}, _update_annotations)
    function_ctx.default_return_type.ret_type = target
    return function_ctx.default_return_type


class AttrgetterPlugin(Plugin):
    def get_function_hook(self, name: str, *args, **kwargs):
        if name != "operator.attrgetter":
            return
        return handle_function_call


def plugin(version) -> Plugin:
    assert float(version) >= 0.720
    return AttrgetterPlugin