mypy plugin to replace types

cli ouput

>> mypy --show-traceback tests/str_int.py
final def (a: builtins.str) -> builtins.int
final <class 'mypy.nodes.FuncDef'>
final def (a: builtins.str) -> builtins.int
final <class 'mypy.nodes.FuncDef'>
tests/str_int.py:2: error: Incompatible return value type (got "int", expected "str")  [return-value]

replace.py

from mypy.plugin import (
    Plugin, 
    FunctionSigContext, 
    FunctionLike
)
from mypy.nodes import FuncDef
from typing_extensions import Callable
from types import FunctionType


class CustomPlugin(Plugin):
    def get_function_signature_hook(
        self, fullname: str
    ) -> Callable[[FunctionSigContext], FunctionLike] | None:
        if "almost_id" in fullname:
            return self.replace_str
        return None

    def replace_str(self, ctx: FunctionSigContext) -> FunctionLike:
        api = ctx.api

        new_signature = ctx.default_signature.copy_modified(
            ret_type=api.named_type('builtins.int')
        )

        un_anal = new_signature.definition.unanalyzed_type.copy_modified(
            ret_type = api.named_type('builtins.int')
        )

        new_signature.definition.unanalyzed_type = un_anal

        typ = new_signature.definition.type.copy_modified(
            ret_type = api.named_type('builtins.int')
        )

        new_signature.definition.type = typ

        print("final", new_signature)
        print("final", type(new_signature.definition))

        return new_signature
      
def plugin(version: str) -> Plugin:
    return CustomPlugin #type: ignore

test_src.py

def almost_id(a: str) -> str:
    return int(a)


almost_id("some")
almost_id("6")