56
56
from mypy .util import split_module_names
57
57
from mypy .typevars import fill_typevars
58
58
from mypy .visitor import ExpressionVisitor
59
- from mypy .plugin import Plugin , MethodContext , MethodSigContext , FunctionContext
59
+ from mypy .plugin import (
60
+ Plugin ,
61
+ MethodContext , MethodSigContext ,
62
+ FunctionContext , FunctionSigContext ,
63
+ )
60
64
from mypy .typeops import (
61
65
tuple_fallback , make_simplified_union , true_only , false_only , erase_to_union_or_bound ,
62
66
function_type , callable_type , try_getting_str_literals , custom_special_method ,
@@ -730,12 +734,15 @@ def apply_function_plugin(self,
730
734
callee .arg_names , formal_arg_names ,
731
735
callee .ret_type , formal_arg_exprs , context , self .chk ))
732
736
733
- def apply_method_signature_hook (
737
+ def apply_signature_hook (
734
738
self , callee : FunctionLike , args : List [Expression ],
735
- arg_kinds : List [int ], context : Context ,
736
- arg_names : Optional [Sequence [Optional [str ]]], object_type : Type ,
737
- signature_hook : Callable [[MethodSigContext ], CallableType ]) -> FunctionLike :
738
- """Apply a plugin hook that may infer a more precise signature for a method."""
739
+ arg_kinds : List [int ],
740
+ arg_names : Optional [Sequence [Optional [str ]]],
741
+ hook : Callable [
742
+ [List [List [Expression ]], CallableType ],
743
+ CallableType ,
744
+ ]) -> FunctionLike :
745
+ """Helper to apply a signature hook for either a function or method"""
739
746
if isinstance (callee , CallableType ):
740
747
num_formals = len (callee .arg_kinds )
741
748
formal_to_actual = map_actuals_to_formals (
@@ -746,19 +753,40 @@ def apply_method_signature_hook(
746
753
for formal , actuals in enumerate (formal_to_actual ):
747
754
for actual in actuals :
748
755
formal_arg_exprs [formal ].append (args [actual ])
749
- object_type = get_proper_type (object_type )
750
- return signature_hook (
751
- MethodSigContext (object_type , formal_arg_exprs , callee , context , self .chk ))
756
+ return hook (formal_arg_exprs , callee )
752
757
else :
753
758
assert isinstance (callee , Overloaded )
754
759
items = []
755
760
for item in callee .items ():
756
- adjusted = self .apply_method_signature_hook (
757
- item , args , arg_kinds , context , arg_names , object_type , signature_hook )
761
+ adjusted = self .apply_signature_hook (
762
+ item , args , arg_kinds , arg_names , hook )
758
763
assert isinstance (adjusted , CallableType )
759
764
items .append (adjusted )
760
765
return Overloaded (items )
761
766
767
+ def apply_function_signature_hook (
768
+ self , callee : FunctionLike , args : List [Expression ],
769
+ arg_kinds : List [int ], context : Context ,
770
+ arg_names : Optional [Sequence [Optional [str ]]],
771
+ signature_hook : Callable [[FunctionSigContext ], CallableType ]) -> FunctionLike :
772
+ """Apply a plugin hook that may infer a more precise signature for a function."""
773
+ return self .apply_signature_hook (
774
+ callee , args , arg_kinds , arg_names ,
775
+ (lambda args , sig :
776
+ signature_hook (FunctionSigContext (args , sig , context , self .chk ))))
777
+
778
+ def apply_method_signature_hook (
779
+ self , callee : FunctionLike , args : List [Expression ],
780
+ arg_kinds : List [int ], context : Context ,
781
+ arg_names : Optional [Sequence [Optional [str ]]], object_type : Type ,
782
+ signature_hook : Callable [[MethodSigContext ], CallableType ]) -> FunctionLike :
783
+ """Apply a plugin hook that may infer a more precise signature for a method."""
784
+ pobject_type = get_proper_type (object_type )
785
+ return self .apply_signature_hook (
786
+ callee , args , arg_kinds , arg_names ,
787
+ (lambda args , sig :
788
+ signature_hook (MethodSigContext (pobject_type , args , sig , context , self .chk ))))
789
+
762
790
def transform_callee_type (
763
791
self , callable_name : Optional [str ], callee : Type , args : List [Expression ],
764
792
arg_kinds : List [int ], context : Context ,
@@ -779,13 +807,17 @@ def transform_callee_type(
779
807
(if appropriate) before the signature is passed to check_call.
780
808
"""
781
809
callee = get_proper_type (callee )
782
- if (callable_name is not None
783
- and object_type is not None
784
- and isinstance (callee , FunctionLike )):
785
- signature_hook = self .plugin .get_method_signature_hook (callable_name )
786
- if signature_hook :
787
- return self .apply_method_signature_hook (
788
- callee , args , arg_kinds , context , arg_names , object_type , signature_hook )
810
+ if callable_name is not None and isinstance (callee , FunctionLike ):
811
+ if object_type is not None :
812
+ method_sig_hook = self .plugin .get_method_signature_hook (callable_name )
813
+ if method_sig_hook :
814
+ return self .apply_method_signature_hook (
815
+ callee , args , arg_kinds , context , arg_names , object_type , method_sig_hook )
816
+ else :
817
+ function_sig_hook = self .plugin .get_function_signature_hook (callable_name )
818
+ if function_sig_hook :
819
+ return self .apply_function_signature_hook (
820
+ callee , args , arg_kinds , context , arg_names , function_sig_hook )
789
821
790
822
return callee
791
823
0 commit comments