@@ -137,6 +137,20 @@ def partial_new_callback(ctx: mypy.plugin.FunctionContext) -> Type:
137
137
# We must normalize from the start to have coherent view together with TypeChecker.
138
138
fn_type = fn_type .with_unpacked_kwargs ().with_normalized_var_args ()
139
139
140
+ last_context = ctx .api .type_context [- 1 ]
141
+ if not fn_type .is_type_obj ():
142
+ # We wrap the return type to get use of a possible type context provided by caller.
143
+ # We cannot do this in case of class objects, since otherwise the plugin may get
144
+ # falsely triggered when evaluating the constructed call itself.
145
+ ret_type : Type = ctx .api .named_generic_type (PARTIAL , [fn_type .ret_type ])
146
+ wrapped_return = True
147
+ else :
148
+ ret_type = fn_type .ret_type
149
+ # Instead, for class objects we ignore any type context to avoid spurious errors,
150
+ # since the type context will be partial[X] etc., not X.
151
+ ctx .api .type_context [- 1 ] = None
152
+ wrapped_return = False
153
+
140
154
defaulted = fn_type .copy_modified (
141
155
arg_kinds = [
142
156
(
@@ -146,7 +160,7 @@ def partial_new_callback(ctx: mypy.plugin.FunctionContext) -> Type:
146
160
)
147
161
for k in fn_type .arg_kinds
148
162
],
149
- ret_type = ctx . api . named_generic_type ( PARTIAL , [ fn_type . ret_type ]) ,
163
+ ret_type = ret_type ,
150
164
)
151
165
if defaulted .line < 0 :
152
166
# Make up a line number if we don't have one
@@ -189,16 +203,20 @@ def partial_new_callback(ctx: mypy.plugin.FunctionContext) -> Type:
189
203
arg_names = actual_arg_names ,
190
204
context = call_expr ,
191
205
)
206
+ if not wrapped_return :
207
+ # Restore previously ignored context.
208
+ ctx .api .type_context [- 1 ] = last_context
209
+
192
210
bound = get_proper_type (bound )
193
211
if not isinstance (bound , CallableType ):
194
212
return ctx .default_return_type
195
- wrapped_ret_type = get_proper_type ( bound . ret_type )
196
- if not isinstance ( wrapped_ret_type , Instance ) or wrapped_ret_type . type . fullname != PARTIAL :
197
- return ctx . default_return_type
198
- if not mypy . semanal . refers_to_fullname ( ctx . args [ 0 ][ 0 ], PARTIAL ):
199
- # If the first argument is partial, above call will trigger the plugin
200
- # again, in between the wrapping above an unwrapping here.
201
- bound = bound .copy_modified (ret_type = wrapped_ret_type .args [0 ])
213
+
214
+ if wrapped_return :
215
+ # Reverse the wrapping we did above.
216
+ ret_type = get_proper_type ( bound . ret_type )
217
+ if not isinstance ( ret_type , Instance ) or ret_type . type . fullname != PARTIAL :
218
+ return ctx . default_return_type
219
+ bound = bound .copy_modified (ret_type = ret_type .args [0 ])
202
220
203
221
formal_to_actual = map_actuals_to_formals (
204
222
actual_kinds = actual_arg_kinds ,
0 commit comments