diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py index 44855f49afaf9..f56535c0954c5 100644 --- a/mypy/checkexpr.py +++ b/mypy/checkexpr.py @@ -3275,15 +3275,23 @@ def real_union(self, typ: Type) -> bool: def type_overrides_set( self, exprs: Sequence[Expression], overrides: Sequence[Type] ) -> Iterator[None]: - """Set _temporary_ type overrides for given expressions.""" + """Set _temporary_ type overrides for given expressions. + + Nested calls are safe: previous overrides (if any) are restored on exit. + """ assert len(exprs) == len(overrides) + previous: list[tuple[Expression, Type | None]] = [] for expr, typ in zip(exprs, overrides): + previous.append((expr, self.type_overrides.get(expr))) self.type_overrides[expr] = typ try: yield finally: - for expr in exprs: - del self.type_overrides[expr] + for expr, old in previous: + if old is None: + self.type_overrides.pop(expr, None) + else: + self.type_overrides[expr] = old def combine_function_signatures(self, types: list[ProperType]) -> AnyType | CallableType: """Accepts a list of function signatures and attempts to combine them together into a @@ -3608,6 +3616,12 @@ def visit_op_expr(self, e: OpExpr) -> Type: if e.op == "*" and isinstance(e.left, ListExpr): # Expressions of form [...] * e get special type inference. return self.check_list_multiply(e) + if e.op == "+": + # Expressions of form list + list under a list type context get special + # type inference so literals/comprehensions honor the outer context. + ctx = get_proper_type(self.type_context[-1]) + if ctx is not None and is_named_instance(ctx, "builtins.list"): + return self.check_list_add(e) if e.op == "%": if isinstance(e.left, BytesExpr): return self.strfrm_checker.check_str_interpolation(e.left, e.right) @@ -4505,6 +4519,32 @@ def check_list_multiply(self, e: OpExpr) -> Type: e.method_type = method_type return result + def check_list_add(self, e: OpExpr) -> Type: + """Type check list concatenation under an outer list type context. + + Like list literals and '[...] * n', constructing operands should see the + outer list[...] context so e.g. ``x: list[int | None] = [0] + [1]`` works. + + Only list literals/comprehensions/nested ``+`` on the right need a type + override; named values should keep their inferred types so ``list.__add__`` + overloads continue to apply (and so we don't nest overrides unsafely). + """ + ctx = self.type_context[-1] + left_type = self.accept(e.left, type_context=ctx) + if isinstance(e.right, (ListExpr, ListComprehension, OpExpr)): + right_type = self.accept(e.right, type_context=ctx) + # check_op re-accepts the right operand; keep the context-aware type. + with self.type_overrides_set([e.right], [right_type]): + result, method_type = self.check_op( + "__add__", left_type, e.right, e, allow_reverse=True + ) + else: + result, method_type = self.check_op( + "__add__", left_type, e.right, e, allow_reverse=True + ) + e.method_type = method_type + return result + def visit_assignment_expr(self, e: AssignmentExpr) -> Type: value = self.accept(e.value) binder_version = self.chk.binder.version diff --git a/test-data/unit/check-inference-context.test b/test-data/unit/check-inference-context.test index 3735f0183eb4a..acf9c1addf5c4 100644 --- a/test-data/unit/check-inference-context.test +++ b/test-data/unit/check-inference-context.test @@ -849,6 +849,20 @@ if int(): a = [''] * 3 # E: List item 0 has incompatible type "str"; expected "int" [builtins fixtures/list.pyi] +[case testListAddInContext] +from typing import List, Optional +x: List[Optional[int]] +x = [0, 1, 2] + [10] +x = [i for i in [0, 1, 2]] + [10] +x = [0] + [None] +x = [] + [1] +x = [1] + [2] + [3] +a: List[Optional[int]] = [None] +x = a + [10] +y: List[int] +y = [0] + ['x'] # E: List item 0 has incompatible type "str"; expected "int" +[builtins fixtures/list.pyi] + [case testUnionTypeContext] from typing import Union, List, TypeVar T = TypeVar('T')