Preserve types of functions registered with singledispatch

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]