Allow plugin signature hooks to return FunctionLike (#10717)

pranavrajpal · web-flow · commit 87a3503e7c90 · 2021-06-26T00:53:43.000-07:00
Change the type of the get_method_signature_hook and
get_function_signature_hook plugin hooks to allow them to return any
FunctionLike type, instead of just CallableType. This allows plugins to
return Overloaded for the type of a function.
diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py
@@ -747,7 +747,7 @@ def apply_signature_hook(
             arg_names: Optional[Sequence[Optional[str]]],
             hook: Callable[
                 [List[List[Expression]], CallableType],
-                CallableType,
+                FunctionLike,
             ]) -> FunctionLike:
         """Helper to apply a signature hook for either a function or method"""
         if isinstance(callee, CallableType):
@@ -775,7 +775,7 @@ def apply_function_signature_hook(
             self, callee: FunctionLike, args: List[Expression],
             arg_kinds: List[int], context: Context,
             arg_names: Optional[Sequence[Optional[str]]],
-            signature_hook: Callable[[FunctionSigContext], CallableType]) -> FunctionLike:
+            signature_hook: Callable[[FunctionSigContext], FunctionLike]) -> FunctionLike:
         """Apply a plugin hook that may infer a more precise signature for a function."""
         return self.apply_signature_hook(
             callee, args, arg_kinds, arg_names,
@@ -786,7 +786,7 @@ def apply_method_signature_hook(
             self, callee: FunctionLike, args: List[Expression],
             arg_kinds: List[int], context: Context,
             arg_names: Optional[Sequence[Optional[str]]], object_type: Type,
-            signature_hook: Callable[[MethodSigContext], CallableType]) -> FunctionLike:
+            signature_hook: Callable[[MethodSigContext], FunctionLike]) -> FunctionLike:
         """Apply a plugin hook that may infer a more precise signature for a method."""
         pobject_type = get_proper_type(object_type)
         return self.apply_signature_hook(
diff --git a/mypy/plugin.py b/mypy/plugin.py
@@ -127,7 +127,9 @@ class C: pass
     Expression, Context, ClassDef, SymbolTableNode, MypyFile, CallExpr
 )
 from mypy.tvar_scope import TypeVarLikeScope
-from mypy.types import Type, Instance, CallableType, TypeList, UnboundType, ProperType
+from mypy.types import (
+    Type, Instance, CallableType, TypeList, UnboundType, ProperType, FunctionLike
+)
 from mypy.messages import MessageBuilder
 from mypy.options import Options
 from mypy.lookup import lookup_fully_qualified
@@ -544,7 +546,7 @@ def func(x: Other[int]) -> None:
         return None
 
     def get_function_signature_hook(self, fullname: str
-                                    ) -> Optional[Callable[[FunctionSigContext], CallableType]]:
+                                    ) -> Optional[Callable[[FunctionSigContext], FunctionLike]]:
         """Adjust the signature of a function.
 
         This method is called before type checking a function call. Plugin
@@ -577,7 +579,7 @@ def get_function_hook(self, fullname: str
         return None
 
     def get_method_signature_hook(self, fullname: str
-                                  ) -> Optional[Callable[[MethodSigContext], CallableType]]:
+                                  ) -> Optional[Callable[[MethodSigContext], FunctionLike]]:
         """Adjust the signature of a method.
 
         This method is called before type checking a method call. Plugin
@@ -748,15 +750,15 @@ def get_type_analyze_hook(self, fullname: str
         return self._find_hook(lambda plugin: plugin.get_type_analyze_hook(fullname))
 
     def get_function_signature_hook(self, fullname: str
-                                    ) -> Optional[Callable[[FunctionSigContext], CallableType]]:
+                                    ) -> Optional[Callable[[FunctionSigContext], FunctionLike]]:
         return self._find_hook(lambda plugin: plugin.get_function_signature_hook(fullname))
 
     def get_function_hook(self, fullname: str
                           ) -> Optional[Callable[[FunctionContext], Type]]:
         return self._find_hook(lambda plugin: plugin.get_function_hook(fullname))
 
     def get_method_signature_hook(self, fullname: str
-                                  ) -> Optional[Callable[[MethodSigContext], CallableType]]:
+                                  ) -> Optional[Callable[[MethodSigContext], FunctionLike]]:
         return self._find_hook(lambda plugin: plugin.get_method_signature_hook(fullname))
 
     def get_method_hook(self, fullname: str
