Skip to content

Commit 902ebb2

Browse files
authored
Use both sides of type context for generic type inference (part 1) (#22286)
Ref #22063 Fixes #4150 Fixes #5311 Fixes #5874 Fixes #5971 Fixes #8861 Fixes #10105 Fixes #12760 Fixes #15087 Fixes #15462 Fixes #15569 Fixes #16310 Fixes #16376 Fixes #16953 Fixes #17025 Fixes #17694 Fixes #18233 Fixes #18236 Fixes #18267 Fixes #18427 Fixes #18689 Fixes #19259 Fixes #19304 Fixes #20219 Fixes #20648 Fixes #21532 Fixes #21796 Fixes #21809 I think this is the right thing to do, even though this unfortunately (but unsurprisingly) causes some visible slow-down (because we infer 2-3x more constraints): ``` master 4.013s (0.0%) | stdev 0.011s new-inference 4.081s (+1.7%) | stdev 0.005s ``` I add couple micro-optimizations in relevant code paths to minimize the impact. The implementation is generally straightforward: use both sides buy default, and only fall back to outer context if there are problems (to give consistent type errors). My plan is to split the change into few smaller things: 1. Basic implementation of the idea (this PR). 2. Targeted improvements of meet/join (as needed). 3. Targeted improvements to the solver such us meet sorting and/or priming. 4. Update the second pass of argument inference in the same manner. 5. Remove (most of) existing special-casing (depends on step 4). 6. Fix false-negative bug due to `Any` leaking from partial context. This is present on master and quite benign, but may happen more often with the new logic. 7. Implement some targeted performance optimizations (like removing `pre_validate_solutions` hack, and making protocol checks faster).
1 parent 49d0cfc commit 902ebb2

14 files changed

Lines changed: 286 additions & 24 deletions

‎mypy/checkexpr.py‎

Lines changed: 81 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
from mypy.checkmember import analyze_member_access, has_operator
2020
from mypy.checkstrformat import StringFormatterChecker
2121
from mypy.constant_fold import constant_fold_expr
22+
from mypy.constraints import Constraint
2223
from mypy.erasetype import erase_type, remove_instance_last_known_values, replace_meta_vars
2324
from mypy.errors import ErrorInfo, ErrorWatcher, report_internal_error
2425
from mypy.expandtype import (
@@ -1792,8 +1793,18 @@ def check_callable_call(
17921793
callee = freshen_function_type_vars(callee)
17931794
self.freshen_cache[(context, original_callee)] = callee
17941795

1796+
initial_callee = callee
1797+
initial_formal_to_actual = None
1798+
constraints: list[Constraint] = []
17951799
if callee.is_generic():
1796-
callee = self.infer_function_type_arguments_using_context(callee, context)
1800+
initial_formal_to_actual = map_actuals_to_formals(
1801+
arg_kinds,
1802+
arg_names,
1803+
callee.arg_kinds,
1804+
callee.arg_names,
1805+
lambda i: self.accept(args[i]),
1806+
)
1807+
callee, constraints = self.infer_function_type_arguments_using_context(callee, context)
17971808

17981809
formal_to_actual = map_actuals_to_formals(
17991810
arg_kinds,
@@ -1825,12 +1836,22 @@ def check_callable_call(
18251836
)
18261837
callee = callee.copy_modified(arg_types=[new_arg_type])
18271838

1828-
if callee.is_generic():
1839+
if initial_callee.is_generic():
18291840
need_refresh = any(
1830-
isinstance(v, (ParamSpecType, TypeVarTupleType)) for v in callee.variables
1841+
isinstance(v, (ParamSpecType, TypeVarTupleType)) for v in initial_callee.variables
18311842
)
1843+
assert initial_formal_to_actual is not None
18321844
callee = self.infer_function_type_arguments(
1833-
callee, args, arg_kinds, arg_names, formal_to_actual, need_refresh, context
1845+
callee,
1846+
initial_callee,
1847+
constraints,
1848+
args,
1849+
arg_kinds,
1850+
arg_names,
1851+
formal_to_actual,
1852+
initial_formal_to_actual,
1853+
need_refresh,
1854+
context,
18341855
)
18351856
if need_refresh:
18361857
# Argument kinds etc. may have changed due to
@@ -2113,7 +2134,7 @@ def infer_arg_types_in_context(
21132134

21142135
def infer_function_type_arguments_using_context(
21152136
self, callable: CallableType, error_context: Context
2116-
) -> CallableType:
2137+
) -> tuple[CallableType, list[Constraint]]:
21172138
"""Unify callable return type to type context to infer type vars.
21182139
21192140
For example, if the return type is set[t] where 't' is a type variable
@@ -2122,7 +2143,7 @@ def infer_function_type_arguments_using_context(
21222143
"""
21232144
ctx = self.type_context[-1]
21242145
if not ctx:
2125-
return callable
2146+
return callable, []
21262147
# The return type may have references to type metavariables that
21272148
# we are inferring right now. We must consider them as indeterminate
21282149
# and they are not potential results; thus we replace them with the
@@ -2179,8 +2200,8 @@ def infer_function_type_arguments_using_context(
21792200
# TODO: we may want to add similar exception if all arguments are lambdas, since
21802201
# in this case external context is almost everything we have.
21812202
if not is_generic_instance(ctx) and not is_literal_type_like(ctx):
2182-
return callable.copy_modified()
2183-
args = infer_type_arguments(
2203+
return callable.copy_modified(), []
2204+
args, constraints = infer_type_arguments(
21842205
callable.variables, ret_type, erased_ctx, skip_unsatisfied=True
21852206
)
21862207
# Only substitute non-Uninhabited and non-erased types.
@@ -2192,23 +2213,28 @@ def infer_function_type_arguments_using_context(
21922213
new_args.append(arg)
21932214
# Don't show errors after we have only used the outer context for inference.
21942215
# We will use argument context to infer more variables.
2195-
return self.apply_generic_arguments(
2196-
callable, new_args, error_context, skip_unsatisfied=True
2216+
return (
2217+
self.apply_generic_arguments(callable, new_args, error_context, skip_unsatisfied=True),
2218+
constraints,
21972219
)
21982220

21992221
def infer_function_type_arguments(
22002222
self,
22012223
callee_type: CallableType,
2224+
initial_callee_type: CallableType,
2225+
ret_constraints: list[Constraint],
22022226
args: list[Expression],
22032227
arg_kinds: list[ArgKind],
22042228
arg_names: Sequence[str | None] | None,
22052229
formal_to_actual: list[list[int]],
2230+
initial_formal_to_actual: list[list[int]],
22062231
need_refresh: bool,
22072232
context: Context,
22082233
) -> CallableType:
22092234
"""Infer the type arguments for a generic callee type.
22102235
2211-
Infer based on the types of arguments.
2236+
Infer based on the types of arguments. The `initial_callee_type` is callee
2237+
type before applying arguments inferred from return type context.
22122238
22132239
Return a derived callable type that has the arguments applied.
22142240
"""
@@ -2233,6 +2259,44 @@ def infer_function_type_arguments(
22332259
else:
22342260
pass1_args.append(arg)
22352261

2262+
# First, try inferring callee type using both argument type and return
2263+
# constraints. Do not do this if special_sig is set or for type objects,
2264+
# those are special-cased elsewhere.
2265+
if (
2266+
initial_callee_type != callee_type
2267+
and initial_callee_type.special_sig is None
2268+
and not initial_callee_type.is_type_obj()
2269+
# This is not very principled, but also prefer outer context for
2270+
# recursive types (because values can never be recursive).
2271+
and not type_state.infer_unions
2272+
):
2273+
inferred_args, _ = infer_function_type_arguments(
2274+
initial_callee_type,
2275+
pass1_args,
2276+
arg_kinds,
2277+
arg_names,
2278+
initial_formal_to_actual,
2279+
context=self.argument_infer_context(),
2280+
strict=self.chk.in_checked_function(),
2281+
extra_constraints=ret_constraints,
2282+
)
2283+
# Only use two-side inference if it is fully successful,
2284+
# otherwise, fallback to return type context as a tiebreaker.
2285+
if not any(
2286+
a is None
2287+
or isinstance(get_proper_type(a), UninhabitedType)
2288+
or has_erased_component(a)
2289+
or set(get_type_vars(a)) & set(callee_type.variables)
2290+
for a in inferred_args
2291+
):
2292+
return self.apply_inferred_arguments(
2293+
initial_callee_type, inferred_args, context
2294+
)
2295+
2296+
if not callee_type.is_generic():
2297+
# Fast path: return context inferred all type arguments.
2298+
return callee_type
2299+
22362300
inferred_args, _ = infer_function_type_arguments(
22372301
callee_type,
22382302
pass1_args,
@@ -6281,7 +6345,12 @@ def accept(
62816345
# context, and use enclosing one, see infer_lambda_type_using_context().
62826346
# TODO: consider using cache for more expression kinds.
62836347
elif (
6284-
isinstance(node, (CallExpr, ListExpr, TupleExpr, DictExpr, OpExpr, MemberExpr))
6348+
(
6349+
isinstance(node, (CallExpr, ListExpr, TupleExpr, DictExpr, OpExpr))
6350+
# Only cache non-static access, it may be slow, and may freshen type vars.
6351+
or isinstance(node, MemberExpr)
6352+
and node.node is None
6353+
)
62856354
and not (self.in_lambda_expr or self.chk.current_node_deferred)
62866355
and not self.chk.options.disable_expression_cache
62876356
):

‎mypy/constraints.py‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -311,7 +311,13 @@ def infer_constraints(
311311
for (t, a) in reversed(type_state.inferring)
312312
):
313313
return []
314-
if has_recursive_types(template) or isinstance(get_proper_type(template), Instance):
314+
# This guard matches the one we use for subtyping, it is slightly more risky,
315+
# but it is much faster than has_recursive_types().
316+
if (
317+
isinstance(template, TypeAliasType)
318+
and template.is_recursive
319+
or isinstance(get_proper_type(template), Instance)
320+
):
315321
# This case requires special care because it may cause infinite recursion.
316322
# Note that we include Instances because the may be recursive as str(Sequence[str]).
317323
if not has_type_vars(template):
@@ -529,6 +535,8 @@ def handle_recursive_union(template: UnionType, actual: Type, direction: int) ->
529535
# the union in two parts, and try inferring sequentially.
530536
non_type_var_items = [t for t in template.items if not isinstance(t, TypeVarType)]
531537
type_var_items = [t for t in template.items if isinstance(t, TypeVarType)]
538+
if not type_var_items or not non_type_var_items:
539+
return []
532540
return infer_constraints(
533541
UnionType.make_union(non_type_var_items), actual, direction
534542
) or infer_constraints(UnionType.make_union(type_var_items), actual, direction)

‎mypy/infer.py‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
from mypy.constraints import (
99
SUBTYPE_OF,
1010
SUPERTYPE_OF,
11+
Constraint,
1112
infer_constraints,
1213
infer_constraints_for_callable,
1314
)
@@ -40,6 +41,7 @@ def infer_function_type_arguments(
4041
context: ArgumentInferContext,
4142
strict: bool = True,
4243
allow_polymorphic: bool = False,
44+
extra_constraints: list[Constraint] | None = None,
4345
) -> tuple[list[Type | None], list[TypeVarLikeType]]:
4446
"""Infer the type arguments of a generic function.
4547
@@ -58,6 +60,8 @@ def infer_function_type_arguments(
5860
constraints = infer_constraints_for_callable(
5961
callee_type, arg_types, arg_kinds, arg_names, formal_to_actual, context
6062
)
63+
if extra_constraints:
64+
constraints += extra_constraints
6165

6266
# Solve constraints.
6367
type_vars = callee_type.variables
@@ -71,10 +75,11 @@ def infer_type_arguments(
7175
is_supertype: bool = False,
7276
skip_unsatisfied: bool = False,
7377
erase_types: bool = True,
74-
) -> list[Type | None]:
78+
) -> tuple[list[Type | None], list[Constraint]]:
7579
# Like infer_function_type_arguments, but only match a single type
7680
# against a generic type.
7781
constraints = infer_constraints(
7882
template, actual, SUPERTYPE_OF if is_supertype else SUBTYPE_OF, erase_types=erase_types
7983
)
80-
return solve_constraints(type_vars, constraints, skip_unsatisfied=skip_unsatisfied)[0]
84+
solutions, _ = solve_constraints(type_vars, constraints, skip_unsatisfied=skip_unsatisfied)
85+
return solutions, constraints

‎mypy/join.py‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -829,8 +829,11 @@ def join_similar_callables(t: CallableType, s: CallableType) -> CallableType:
829829
# The result should always use 'function' as a fallback if either operands are using it.
830830
if t.fallback.type.fullname == "builtins.function":
831831
fallback = t.fallback
832-
else:
832+
elif s.fallback.type.fullname == "builtins.function":
833833
fallback = s.fallback
834+
else:
835+
fallback = InstanceJoiner().join_instances(t.fallback, s.fallback)
836+
assert isinstance(fallback, Instance)
834837
instance_type = None
835838
if t.instance_type is not None and s.instance_type is not None:
836839
instance_type = join_types(t.instance_type, s.instance_type)
@@ -888,8 +891,11 @@ def combine_similar_callables(t: CallableType, s: CallableType) -> CallableType:
888891
# The result should always use 'function' as a fallback if either operands are using it.
889892
if t.fallback.type.fullname == "builtins.function":
890893
fallback = t.fallback
891-
else:
894+
elif s.fallback.type.fullname == "builtins.function":
892895
fallback = s.fallback
896+
else:
897+
fallback = InstanceJoiner().join_instances(t.fallback, s.fallback)
898+
assert isinstance(fallback, Instance)
893899
instance_type = None
894900
if t.instance_type is not None and s.instance_type is not None:
895901
instance_type = join_types(t.instance_type, s.instance_type)

‎mypy/messages.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3136,7 +3136,7 @@ def [T <: int] f(self, x: int, y: T) -> None
31363136
# TODO: avoid different logic for incremental runs.
31373137
first_arg = get_first_arg(tp)
31383138

3139-
if tp.is_type_obj():
3139+
if tp.is_type_obj() or not tp.is_bound:
31403140
skip_self = True
31413141
if first_arg and not skip_self:
31423142
if s:

‎mypy/typeops.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -561,7 +561,7 @@ def bind_self_inner(
561561
return type_state.bind_self_cache[cache_key]
562562

563563
# Solve for these type arguments using the actual class or instance type.
564-
typeargs = infer_type_arguments(
564+
typeargs, _ = infer_type_arguments(
565565
self_vars, self_param_type, original_type, is_supertype=True, erase_types=False
566566
)
567567
if (
@@ -570,7 +570,7 @@ def bind_self_inner(
570570
and isinstance(original_type, (Instance, TypeVarType, TupleType))
571571
):
572572
# In case we call a classmethod through an instance x, fallback to type(x).
573-
typeargs = infer_type_arguments(
573+
typeargs, _ = infer_type_arguments(
574574
self_vars,
575575
self_param_type,
576576
TypeType(original_type),

‎test-data/unit/check-functions.test‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3957,3 +3957,25 @@ Ex = TypeVar("Ex", covariant=True)
39573957

39583958
def takes_lambda(func: Callable[[Ex], Any]) -> None:
39593959
pass
3960+
3961+
[case testTypeObjectJoinWithOverloadsCorrectFallback]
3962+
from typing import Iterable, overload
3963+
3964+
class M(type): ...
3965+
class A(metaclass=M): ...
3966+
class B: ...
3967+
class C(metaclass=M):
3968+
@overload
3969+
def __init__(self, x: int) -> None: ...
3970+
@overload
3971+
def __init__(self, x: str) -> None: ...
3972+
def __init__(self, x):
3973+
pass
3974+
3975+
def f(x: Iterable[object]) -> None: ...
3976+
f([A, B, C])
3977+
f([B, A, C])
3978+
f([B, C, A])
3979+
f([A, C, B])
3980+
f([C, A, B])
3981+
f([C, B, A])

‎test-data/unit/check-inference-context.test‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1635,3 +1635,48 @@ def foo(args: Union[Iterable[Union[str, int]], str, int]) -> Iterable[str]:
16351635
)
16361636
return args
16371637
[builtins fixtures/isinstancelist.pyi]
1638+
1639+
[case testInferenceContextBothSidesInvariant]
1640+
from typing import TypeVar, Generic
1641+
1642+
T = TypeVar("T")
1643+
1644+
class B(Generic[T]): ...
1645+
class C(B[int]): ...
1646+
1647+
def first(xs: list[T]) -> T: ...
1648+
1649+
b: B[int]
1650+
cs: list[C]
1651+
b = first(cs) # OK
1652+
reveal_type(b) # N: Revealed type is "__main__.C"
1653+
1654+
[case testInferenceContextBothSidesNarrow]
1655+
from typing import TypeVar, Generic
1656+
1657+
T = TypeVar("T")
1658+
1659+
class B(Generic[T]): ...
1660+
class C(B[int]): ...
1661+
1662+
def same(x: T) -> T: ...
1663+
1664+
b: B[int]
1665+
b = same(C()) # This should narrow
1666+
reveal_type(b) # N: Revealed type is "__main__.C"
1667+
1668+
[case testInferenceContextBothSidesJoin]
1669+
from typing import TypeVar, Generic, Callable, Union
1670+
1671+
T = TypeVar("T")
1672+
1673+
class A: ...
1674+
class B: ...
1675+
class C: ...
1676+
1677+
def make(x: T) -> Callable[[T], None]: ...
1678+
1679+
x: Union[B, C]
1680+
fn: Callable[[Union[A, B]], None]
1681+
fn = make(x)
1682+
reveal_type(fn) # N: Revealed type is "def (__main__.B | __main__.C | __main__.A)"

‎test-data/unit/check-parameter-specification.test‎

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2814,6 +2814,20 @@ x: Union[Any, C[[Any, Any]]]
28142814

28152815
def test(x: int) -> None: ...
28162816

2817-
x = foo(test) # E: Argument 1 to "foo" has incompatible type "Callable[[int], None]"; expected "Callable[[Any, Any], None]"
2818-
reveal_type(x) # N: Revealed type is "__main__.C[[Any, Any]]"
2817+
x = foo(test)
2818+
reveal_type(x) # N: Revealed type is "__main__.C[Any] | __main__.C[[Any, Any]]"
2819+
[builtins fixtures/paramspec.pyi]
2820+
2821+
[case testTwoSidedInferenceWithParamSpec]
2822+
from typing import Callable, Generic, TypeVar
2823+
from typing_extensions import ParamSpec
2824+
2825+
P = ParamSpec("P")
2826+
T = TypeVar("T")
2827+
class C(Generic[P, T]): ...
2828+
2829+
def thing(func: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> C[P, T]: ...
2830+
2831+
def test(x: int = 0) -> None: ...
2832+
f: C[[], None] = thing(test)
28192833
[builtins fixtures/paramspec.pyi]

0 commit comments

Comments
 (0)