Preserve types of functions registered with singledispatch (#10756)

When calling a function that has been registered as an implementation
for a singledispatch function, currently, mypy doesn't type check
arguments to those functions when those functions are called directly
(instead of through the main singledispatch function). This fixes that by
preserving the type of the registered function even after the register decorator
has been used.

Author: pranavrajpal

mypy/plugins/singledispatch.py

@@ -139,13 +139,19 @@ def singledispatch_register_callback(ctx: MethodContext) -> Type:
     elif isinstance(first_arg_type, CallableType):
         # TODO: do more checking for registered functions
         register_function(ctx, ctx.type, first_arg_type)
+        # The typeshed stubs for register say that the function returned is Callable[..., T], even
+        # though the function returned is the same as the one passed in. We return the type of the
+        # function so that mypy can properly type check cases where the registered function is used
+        # directly (instead of through singledispatch)
+        return first_arg_type
 
-    # register doesn't modify the function it's used on
+    # fallback in case we don't recognize the arguments
     return ctx.default_return_type
 
 
 def register_function(ctx: PluginContext, singledispatch_obj: Instance, func: Type,
                       register_arg: Optional[Type] = None) -> None:
+    """Register a function"""
 
     func = get_proper_type(func)
     if not isinstance(func, CallableType):
@@ -165,6 +171,7 @@ def register_function(ctx: PluginContext, singledispatch_obj: Instance, func: Ty
                 format_type(dispatch_type), format_type(fallback_dispatch_type)
             ), func.definition)
         return
+    return
 
 
 def get_dispatch_type(func: CallableType, register_arg: Optional[Type]) -> Optional[Type]:
@@ -183,6 +190,8 @@ def call_singledispatch_function_after_register_argument(ctx: MethodContext) ->
         func = get_first_arg(ctx.arg_types)
         if func is not None:
             register_function(ctx, type_args.singledispatch_obj, func, type_args.register_type)
+            # see call to register_function in the callback for register
+            return func
     return ctx.default_return_type
 
 

test-data/unit/check-singledispatch.test

@@ -289,3 +289,28 @@ def h(a: Missing) -> None:
     pass
 
 [builtins fixtures/args.pyi]
+
+[case testIncorrectArgumentTypeWhenCallingRegisteredImplDirectly]
+from functools import singledispatch
+
+@singledispatch
+def f(arg, arg2: str) -> bool:
+    return False
+
+@f.register
+def g(arg: int, arg2: str) -> bool:
+    pass
+
+@f.register(str)
+def h(arg, arg2: str) -> bool:
+    pass
+
+g('a', 'a') # E: Argument 1 to "g" has incompatible type "str"; expected "int"
+g(1, 1) # E: Argument 2 to "g" has incompatible type "int"; expected "str"
+
+# don't show errors for incorrect first argument here, because there's no type annotation for the
+# first argument
+h(1, 'a')
+h('a', 1) # E: Argument 2 to "h" has incompatible type "int"; expected "str"
+
+[builtins fixtures/args.pyi]