Skip to content

Commit a0dda1c

Browse files
ilevkivskyiilevkivskyi
authored
Use both sides of type context for generic type inference (part 2) (#22394)
Ref #22063 Fixes #8478 Fixes #10281 This makes the new logic work in presence of second pass arguments (callables). I also did a little refactoring, and added some documenting comments. Implementation is generally straightforward.
1 parent 6349140 commit a0dda1c

4 files changed

Lines changed: 151 additions & 18 deletions

File tree

‎mypy/checkexpr.py‎

Lines changed: 118 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +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
22+
from mypy.constraints import Constraint, infer_constraints_for_callable
2323
from mypy.erasetype import erase_type, remove_instance_last_known_values, replace_meta_vars
2424
from mypy.errors import ErrorInfo, ErrorWatcher, report_internal_error
2525
from mypy.expandtype import (
@@ -121,6 +121,7 @@
121121
Plugin,
122122
)
123123
from mypy.semanal_enum import ENUM_BASES
124+
from mypy.solve import solve_constraints
124125
from mypy.state import state
125126
from mypy.subtypes import (
126127
common_type,
@@ -1796,6 +1797,21 @@ def check_callable_call(
17961797
initial_callee = callee
17971798
initial_formal_to_actual = None
17981799
constraints: list[Constraint] = []
1800+
# The general flow of type inference is the following:
1801+
# * Infer constraints from outer (return) context and solve them.
1802+
# * Apply the inferred type arguments to the callable.
1803+
# * Infer argument types in this partial context.
1804+
# * Infer constraints from inner (argument) context (pass 1, non-callable arguments).
1805+
# * Solve the combined constraint set (outer + partial inner).
1806+
# * Apply the new (more precise) inferred type arguments to initial callable.
1807+
# * Re-infer argument types in more precise type context
1808+
# (including pass 2, callable arguments).
1809+
# * Re-infer constraints from inner context (now from all arguments).
1810+
# * Solve the final constraint set combining outer context and precise inner context.
1811+
# Apply the inferred final type arguments to the initial generic callable.
1812+
# * If we cannot solve (likely an error), we repeat the process starting from outer
1813+
# context. Now progressively applying inferred type arguments after each step,
1814+
# to give consistent error messages.
17991815
if callee.is_generic():
18001816
initial_formal_to_actual = map_actuals_to_formals(
18011817
arg_kinds,
@@ -2218,6 +2234,84 @@ def infer_function_type_arguments_using_context(
22182234
constraints,
22192235
)
22202236

2237+
def infer_function_type_arguments_both_sides(
2238+
self,
2239+
initial_callee_type: CallableType,
2240+
ret_constraints: list[Constraint],
2241+
args: list[Expression],
2242+
pass1_arg_types: list[Type | None],
2243+
arg_pass_nums: list[int],
2244+
arg_kinds: list[ArgKind],
2245+
arg_names: Sequence[str | None] | None,
2246+
formal_to_actual: list[list[int]],
2247+
initial_formal_to_actual: list[list[int]],
2248+
need_refresh: bool,
2249+
context: Context,
2250+
) -> CallableType | None:
2251+
"""Infer type arguments of a generic callable by combining outer and inner contexts."""
2252+
# Inline infer_function_type_arguments() to get access to argument-side constraints.
2253+
constraints = infer_constraints_for_callable(
2254+
initial_callee_type,
2255+
pass1_arg_types,
2256+
arg_kinds,
2257+
arg_names,
2258+
initial_formal_to_actual,
2259+
self.argument_infer_context(),
2260+
)
2261+
constraints += ret_constraints
2262+
2263+
type_vars = initial_callee_type.variables
2264+
type_vars_set = set(type_vars)
2265+
inferred_args, _ = solve_constraints(
2266+
type_vars, constraints, strict=self.chk.in_checked_function()
2267+
)
2268+
2269+
pass1_inferred_args = inferred_args
2270+
if 2 in arg_pass_nums:
2271+
# This is an adjusted version of logic in infer_function_type_arguments_pass2().
2272+
# Here we proceed even if first pass inferred all type arguments, since
2273+
# second pass may infer narrower types in presence of wide outer context.
2274+
solid_inferred_args = list(inferred_args)
2275+
for i, arg in enumerate(inferred_args):
2276+
if not solid_inferred_type(arg, type_vars_set):
2277+
solid_inferred_args[i] = None
2278+
pass2_callee_type = self.apply_generic_arguments(
2279+
initial_callee_type, solid_inferred_args, context
2280+
)
2281+
if need_refresh:
2282+
formal_to_actual = map_actuals_to_formals(
2283+
arg_kinds,
2284+
arg_names,
2285+
pass2_callee_type.arg_kinds,
2286+
pass2_callee_type.arg_names,
2287+
lambda a: self.accept(args[a]),
2288+
)
2289+
2290+
# (Re-)infer argument types in the more precise type context (for lambdas).
2291+
with self.msg.filter_errors():
2292+
arg_types = self.infer_arg_types_in_context(
2293+
pass2_callee_type, args, arg_kinds, formal_to_actual
2294+
)
2295+
inferred_args, _ = infer_function_type_arguments(
2296+
initial_callee_type,
2297+
arg_types,
2298+
arg_kinds,
2299+
arg_names,
2300+
initial_formal_to_actual,
2301+
context=self.argument_infer_context(),
2302+
extra_constraints=constraints,
2303+
)
2304+
2305+
# Only use two-side inference if it is fully successful,
2306+
# otherwise, fallback to return type context as a tiebreaker.
2307+
if all(solid_inferred_type(a, type_vars_set) for a in inferred_args):
2308+
return self.apply_inferred_arguments(initial_callee_type, inferred_args, context)
2309+
# Before giving up, check first pass result. In some cases this may help in situations
2310+
# where both passes together create a system too tricky for our solver.
2311+
if all(solid_inferred_type(a, type_vars_set) for a in pass1_inferred_args):
2312+
return self.apply_inferred_arguments(initial_callee_type, pass1_inferred_args, context)
2313+
return None
2314+
22212315
def infer_function_type_arguments(
22222316
self,
22232317
callee_type: CallableType,
@@ -2270,28 +2364,21 @@ def infer_function_type_arguments(
22702364
# recursive types (because values can never be recursive).
22712365
and not type_state.infer_unions
22722366
):
2273-
inferred_args, _ = infer_function_type_arguments(
2367+
inferred_callee = self.infer_function_type_arguments_both_sides(
22742368
initial_callee_type,
2369+
ret_constraints,
2370+
args,
22752371
pass1_args,
2372+
arg_pass_nums,
22762373
arg_kinds,
22772374
arg_names,
2375+
formal_to_actual,
22782376
initial_formal_to_actual,
2279-
context=self.argument_infer_context(),
2280-
strict=self.chk.in_checked_function(),
2281-
extra_constraints=ret_constraints,
2377+
need_refresh,
2378+
context,
22822379
)
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-
)
2380+
if inferred_callee is not None:
2381+
return inferred_callee
22952382

22962383
if not callee_type.is_generic():
22972384
# Fast path: return context inferred all type arguments.
@@ -6904,6 +6991,20 @@ def visit_uninhabited_type(self, t: UninhabitedType) -> bool:
69046991
return t.ambiguous
69056992

69066993

6994+
def solid_inferred_type(typ: Type | None, type_vars: set[TypeVarLikeType]) -> bool:
6995+
"""Are we sure this is a good inferred type argument for a generic callable?"""
6996+
if typ is None:
6997+
return False
6998+
if isinstance(get_proper_type(typ), UninhabitedType):
6999+
return False
7000+
if has_erased_component(typ):
7001+
return False
7002+
# Make sure there are no type variable leaks in the type.
7003+
if set(get_type_vars(typ)) & type_vars:
7004+
return False
7005+
return True
7006+
7007+
69077008
def arg_approximate_similarity(actual: Type, formal: Type) -> bool:
69087009
"""Return if caller argument (actual) is roughly compatible with signature arg (formal).
69097010

‎mypy/infer.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,12 +9,13 @@
99
SUBTYPE_OF,
1010
SUPERTYPE_OF,
1111
Constraint,
12+
filter_imprecise_kinds,
1213
infer_constraints,
1314
infer_constraints_for_callable,
1415
)
1516
from mypy.nodes import ArgKind
1617
from mypy.solve import solve_constraints
17-
from mypy.types import CallableType, Instance, Type, TypeVarLikeType
18+
from mypy.types import CallableType, Instance, ParamSpecType, Type, TypeVarLikeType
1819

1920

2021
class ArgumentInferContext:
@@ -62,6 +63,9 @@ def infer_function_type_arguments(
6263
)
6364
if extra_constraints:
6465
constraints += extra_constraints
66+
# Combining extra constraints may shadow precise constraints by imprecise ones.
67+
if any(isinstance(v, ParamSpecType) for v in callee_type.variables):
68+
constraints = filter_imprecise_kinds(constraints)
6569

6670
# Solve constraints.
6771
type_vars = callee_type.variables

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1680,3 +1680,14 @@ x: Union[B, C]
16801680
fn: Callable[[Union[A, B]], None]
16811681
fn = make(x)
16821682
reveal_type(fn) # N: Revealed type is "def (__main__.B | __main__.C | __main__.A)"
1683+
1684+
[case testInferenceContextBothSidesCallable]
1685+
from typing import TypeVar, Callable, Optional
1686+
1687+
T = TypeVar("T")
1688+
S = TypeVar("S")
1689+
1690+
def run(x: list[T], fn: Callable[[T], S]) -> list[Optional[S]]: ...
1691+
1692+
def worker(d: int) -> Optional[int]: ...
1693+
res: list[Optional[int]] = run([1, 2, 3], worker) # OK

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

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2831,3 +2831,20 @@ def thing(func: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> C[P, T]: .
28312831
def test(x: int = 0) -> None: ...
28322832
f: C[[], None] = thing(test)
28332833
[builtins fixtures/paramspec.pyi]
2834+
2835+
[case testTwoSidedInferenceWithParamSpecGeneric]
2836+
from typing import Callable, Generic, TypeVar, Union
2837+
from typing_extensions import ParamSpec
2838+
2839+
P = ParamSpec("P")
2840+
T = TypeVar("T")
2841+
I = TypeVar("I", bound=Union[int, str])
2842+
2843+
class C(Generic[T]): ...
2844+
2845+
def apply(func: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> T: ...
2846+
2847+
def test(x: C[I]) -> list[int]: ...
2848+
ci: C[int]
2849+
f: list[int] = apply(test, ci)
2850+
[builtins fixtures/paramspec.pyi]

0 commit comments

Comments
 (0)