diff --git a/crates/ty_python_semantic/resources/mdtest/conditional/if_expression.md b/crates/ty_python_semantic/resources/mdtest/conditional/if_expression.md index 428e5dfb94..bbf05c5819 100644 --- a/crates/ty_python_semantic/resources/mdtest/conditional/if_expression.md +++ b/crates/ty_python_semantic/resources/mdtest/conditional/if_expression.md @@ -74,25 +74,78 @@ def _(flag: bool): incompatible_mapping: dict[str, int] = {"x": "y"} if flag else {} # error: [invalid-assignment] ``` -## Other collection branches +## Joint collection branches -When both branches are empty, their element types remain unknown. A list in the other branch can -also supply `Unknown` or `Any` as the element type. Lists with different element types combine those -types. A list in one branch and a dictionary in the other produce a union of the two types. +If both branches contain literals for the same collection type, the type inferred for the +conditional expression combines both literals: + +```py +def _(flag: bool): + a = ["x"] if flag else [1] + reveal_type(a) # revealed: list[str | int] + + b = [1] if flag else ["x"] + reveal_type(b) # revealed: list[int | str] + + items = {"x"} if flag else {1} + reveal_type(items) # revealed: set[str | int] + + mapping = {"key": 1} if flag else {2: "value"} + reveal_type(mapping) # revealed: dict[str | int, int | str] + + _ = ["x"] if flag else [int(missing)] # error: [unresolved-reference] +``` + +Dynamic types are retained from both branches: ```py from typing import Any -from ty_extensions._internal import Unknown -def _(flag: bool, unknown: list[Unknown], gradual: list[Any]): +def _(flag: bool, dynamic: Any): + a = [dynamic] if flag else [1] + reveal_type(a) # revealed: list[Any | int] + + b = [1] if flag else [dynamic] + reveal_type(b) # revealed: list[int | Any] + + nested = [[]] if flag else [[1]] + reveal_type(nested) # revealed: list[list[Unknown] | list[int]] +``` + +When both branches are empty, their element types remain unknown: + +```py +def _(flag: bool): reveal_type([] if flag else []) # revealed: list[Unknown] reveal_type({} if flag else {}) # revealed: dict[Unknown, Unknown] - reveal_type(["x"] if flag else unknown) # revealed: list[Unknown] - reveal_type(["x"] if flag else gradual) # revealed: list[Any] - reveal_type(["x"] if flag else [1]) # revealed: list[int | str] +``` + +A list in one branch and a dictionary in the other produce a union of the two types: + +```py +def _(flag: bool): reveal_type(["x"] if flag else {}) # revealed: list[str] | dict[Unknown, Unknown] ``` +If one branch contains an existing collection, the literal in the other branch is inferred against +the existing collection's declared type: + +```py +def _(flag: bool, existing: list[int], gradual: list[Any]): + reveal_type(["x"] if flag else gradual) # revealed: list[Any] + + aliased = ["x"] if flag else existing + reveal_type(aliased) # revealed: list[int | str] | list[int] +``` + +An explicit union annotation can supply a different element type to each branch: + +```py +def _(flag: bool): + x: list[int] | list[str] = [1] if flag else ["x"] + reveal_type(x) # revealed: list[int] | list[str] +``` + ## Collection branches and evaluation order The condition is evaluated before either branch, and a branch's assignment expression remains @@ -112,6 +165,9 @@ def _(flag: bool): result = [(skipped := "after")] if False else [] reveal_type(result) # revealed: list[str] reveal_type(skipped) # revealed: Literal["before"] + + reachable = ["x"] if False else [1] + reveal_type(reachable) # revealed: list[int] ``` ## Statically known compound conditions diff --git a/crates/ty_python_semantic/src/types/infer/builder.rs b/crates/ty_python_semantic/src/types/infer/builder.rs index 97067ac166..397502d0be 100644 --- a/crates/ty_python_semantic/src/types/infer/builder.rs +++ b/crates/ty_python_semantic/src/types/infer/builder.rs @@ -7498,63 +7498,27 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { } fn infer_list_expression(&mut self, list: &ast::ExprList, tcx: TypeContext<'db>) -> Type<'db> { - let db = self.db(); - let ast::ExprList { - range: _, - node_index: _, - elts, - ctx: _, - } = list; - - let elts = elts.iter().map(|elt| [Some(elt)]).collect_vec(); - let mut infer_elt_ty = - |builder: &mut Self, (_, elt, tcx)| builder.infer_expression(elt, tcx); + let elts = list.elts.iter().map(|elt| [Some(elt)]).collect_vec(); + let infer_elt_ty = |builder: &mut Self, (_, elt, tcx)| builder.infer_expression(elt, tcx); self.infer_collection_literal( KnownClass::List, Some(list.into()), &elts, - &mut infer_elt_ty, + infer_elt_ty, tcx, ) - .unwrap_or_else(|| { - KnownClass::List.to_specialized_instance( - db, - self.program_environment(), - &[Type::unknown()], - ) - }) } fn infer_set_expression(&mut self, set: &ast::ExprSet, tcx: TypeContext<'db>) -> Type<'db> { - let db = self.db(); - let ast::ExprSet { - range: _, - node_index: _, - elts, - } = set; - - let elts = elts.iter().map(|elt| [Some(elt)]).collect_vec(); + let elts = set.elts.iter().map(|elt| [Some(elt)]).collect_vec(); let fallback_tcx = self.incomplete_typed_dict_key_context(set, tcx); - let mut infer_elt_ty = |builder: &mut Self, arg: ArgExpr<'db, '_>| { + let infer_elt_ty = |builder: &mut Self, arg: ArgExpr<'db, '_>| { let (_, elt, elt_tcx) = arg; builder.infer_set_element(elt, elt_tcx, fallback_tcx) }; - self.infer_collection_literal( - KnownClass::Set, - Some(set.into()), - &elts, - &mut infer_elt_ty, - tcx, - ) - .unwrap_or_else(|| { - KnownClass::Set.to_specialized_instance( - db, - self.program_environment(), - &[Type::unknown()], - ) - }) + self.infer_collection_literal(KnownClass::Set, Some(set.into()), &elts, infer_elt_ty, tcx) } /// Infers a set element, optionally with a fallback context for an incomplete `TypedDict` key. @@ -7714,7 +7678,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { // Avoid inferring the items multiple times if we already attempted to infer the // dictionary literal as a `TypedDict`. This also allows us to infer using the // type context of the expected `TypedDict` field. - let mut infer_elt_ty = |builder: &mut Self, (_, elt, tcx): ArgExpr<'db, '_>| { + let infer_elt_ty = |builder: &mut Self, (_, elt, tcx): ArgExpr<'db, '_>| { item_types .get(&elt.node_index().load()) .copied() @@ -7726,16 +7690,31 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { KnownClass::Dict, Some(dict.into()), &items, - &mut infer_elt_ty, + infer_elt_ty, tcx, ) - .unwrap_or_else(|| { - KnownClass::Dict.to_specialized_instance(db, env, &[Type::unknown(), Type::unknown()]) - }) + } + + fn infer_collection_literal<'expr, const N: usize>( + &mut self, + class: KnownClass, + expression: Option>, + elements: &[[Option<&'expr ast::Expr>; N]], + mut infer_element: impl FnMut(&mut Self, ArgExpr<'db, 'expr>) -> Type<'db>, + tcx: TypeContext<'db>, + ) -> Type<'db> { + self.try_infer_collection_literal(class, expression, elements, &mut infer_element, tcx) + .unwrap_or_else(|| { + class.to_specialized_instance( + self.db(), + self.program_environment(), + &[Type::unknown(); N], + ) + }) } // Infer the type of a collection literal expression. - fn infer_collection_literal<'expr, const N: usize>( + fn try_infer_collection_literal<'expr, const N: usize>( &mut self, collection_class: KnownClass, collection_expr: Option>, @@ -7749,7 +7728,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { let mut speculative_builder = self.speculate(); // Attempt to infer the collection literal using the narrowed type context. - let inferred_ty = speculative_builder.infer_collection_literal_impl( + let inferred_ty = speculative_builder.try_infer_collection_literal_impl( collection_class, collection_expr, elts, @@ -7780,7 +7759,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { } } - self.infer_collection_literal_impl( + self.try_infer_collection_literal_impl( collection_class, collection_expr, elts, @@ -7790,7 +7769,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { } // Infer the type of a collection literal expression. - fn infer_collection_literal_impl<'expr, const N: usize>( + fn try_infer_collection_literal_impl<'expr, const N: usize>( &mut self, collection_class: KnownClass, collection_expr: Option>, @@ -8416,15 +8395,12 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { elements: [Option<&ast::Expr>; N], inference: &ScopeInference<'db>, tcx: TypeContext<'db>, - ) -> Option> { - let mut infer_element_ty = - |_builder: &mut Self, (_, elt, _)| inference.expression_type(elt); - + ) -> Type<'db> { self.infer_collection_literal( collection_class, Some(collection_expr), &[elements], - &mut infer_element_ty, + |_, (_, elt, _)| inference.expression_type(elt), tcx, ) } @@ -8434,7 +8410,6 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { listcomp: &ast::ExprListComp, tcx: TypeContext<'db>, ) -> Type<'db> { - let db = self.db(); let ast::ExprListComp { range: _, node_index: _, @@ -8461,13 +8436,6 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { inference, tcx, ) - .unwrap_or_else(|| { - KnownClass::List.to_specialized_instance( - db, - self.program_environment(), - &[Type::unknown()], - ) - }) } fn infer_set_comprehension_expression( @@ -8475,7 +8443,6 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { setcomp: &ast::ExprSetComp, tcx: TypeContext<'db>, ) -> Type<'db> { - let db = self.db(); let ast::ExprSetComp { range: _, node_index: _, @@ -8502,13 +8469,6 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { inference, tcx, ) - .unwrap_or_else(|| { - KnownClass::Set.to_specialized_instance( - db, - self.program_environment(), - &[Type::unknown()], - ) - }) } fn infer_dict_comprehension_expression( @@ -8516,7 +8476,6 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { dictcomp: &ast::ExprDictComp, tcx: TypeContext<'db>, ) -> Type<'db> { - let db = self.db(); let ast::ExprDictComp { range: _, node_index: _, @@ -8544,13 +8503,6 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { inference, tcx, ) - .unwrap_or_else(|| { - KnownClass::Dict.to_specialized_instance( - db, - self.program_environment(), - &[Type::unknown(), Type::unknown()], - ) - }) } fn infer_generator_expression_scope( @@ -8599,7 +8551,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { let mut infer_elt_ty = |builder: &mut Self, (_, elt, tcx)| builder.infer_expression(elt, tcx); - self.infer_collection_literal( + self.try_infer_collection_literal( KnownClass::List, Some(listcomp.into()), &elts, @@ -8627,7 +8579,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { let mut infer_elt_ty = |builder: &mut Self, (_, elt, tcx)| builder.infer_expression(elt, tcx); - self.infer_collection_literal( + self.try_infer_collection_literal( KnownClass::Set, Some(setcomp.into()), &elts, @@ -8657,7 +8609,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { let mut infer_elt_ty = |builder: &mut Self, (_, elt, tcx)| builder.infer_expression(elt, tcx); - self.infer_collection_literal( + self.try_infer_collection_literal( KnownClass::Dict, Some(dictcomp.into()), &elts, @@ -8871,28 +8823,6 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { } = if_expression; let test_ty = self.infer_maybe_standalone_expression(test, TypeContext::default()); - let (body_ty, orelse_ty) = if is_collection_literal(body) - && prefer_collection_literal_peer_context(db, env, tcx) - && !is_nonempty_collection_literal_with_empty_peer(body, orelse) - { - // Infer the peer branch first so the body can use its type as context. - let orelse_ty = self.infer_expression(orelse, tcx); - let body_ty = self.infer_expression_with_collection_literal_peer_context( - body, - tcx, - Some(orelse_ty), - ); - (body_ty, orelse_ty) - } else { - let body_ty = self.infer_expression(body, tcx); - let orelse_ty = self.infer_expression_with_collection_literal_peer_context( - orelse, - tcx, - Some(body_ty), - ); - (body_ty, orelse_ty) - }; - let test_truthiness = match test_ty.try_bool(db, env) { Ok(_) => self.expression_truthiness(test, ExpressionContext::Condition), Err(err) => { @@ -8903,6 +8833,26 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { self.check_condition_redundancy(test, test_ty); + let prefer_peer = + is_collection_literal(body) && prefer_collection_literal_peer_context(db, env, tcx); + if prefer_peer + && test_truthiness.is_ambiguous() + && let Some(ty) = self.infer_joint_collection_literals(body, orelse, tcx) + { + self.finish_expression_type(body, ty, tcx); + self.finish_expression_type(orelse, ty, tcx); + return ty; + } + + // Infer the peer branch first so the body can use its type as context. + let orelse_ty = (prefer_peer && !is_empty_collection_literal(orelse)) + .then(|| self.infer_expression(orelse, tcx)); + let body_ty = + self.infer_expression_with_collection_literal_peer_context(body, tcx, orelse_ty); + let orelse_ty = orelse_ty.unwrap_or_else(|| { + self.infer_expression_with_collection_literal_peer_context(orelse, tcx, Some(body_ty)) + }); + match test_truthiness { Truthiness::AlwaysTrue => body_ty, Truthiness::AlwaysFalse => orelse_ty, @@ -8911,6 +8861,60 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { } } + /// Infer a merged collection type from two distinct literal expressions. + fn infer_joint_collection_literals( + &mut self, + body: &ast::Expr, + orelse: &ast::Expr, + tcx: TypeContext<'db>, + ) -> Option> { + // A generic TypedDict can have known fields even when its type arguments are unsolved. + if tcx + .annotation + .is_some_and(|ty| ty.as_typed_dict().is_some()) + { + return None; + } + + let infer_elt_ty = |builder: &mut Self, (_, elt, tcx): ArgExpr<'db, '_>| { + builder.infer_expression(elt, tcx) + }; + + let (class, left, right) = match (body, orelse) { + (ast::Expr::List(left), ast::Expr::List(right)) => { + (KnownClass::List, &left.elts, &right.elts) + } + (ast::Expr::Set(left), ast::Expr::Set(right)) => { + (KnownClass::Set, &left.elts, &right.elts) + } + (ast::Expr::Dict(left), ast::Expr::Dict(right)) => { + let items = left + .items + .iter() + .chain(&right.items) + .map(|item| [item.key.as_ref(), Some(&item.value)]) + .collect_vec(); + + return Some(self.infer_collection_literal( + KnownClass::Dict, + None, + &items, + infer_elt_ty, + tcx, + )); + } + _ => return None, + }; + + let elts = left + .iter() + .chain(right) + .map(|elt| [Some(elt)]) + .collect_vec(); + + Some(self.infer_collection_literal(class, None, &elts, infer_elt_ty, tcx)) + } + fn infer_lambda_body(&mut self, lambda_expression: &ast::ExprLambda, tcx: TypeContext<'db>) { self.infer_expression(&lambda_expression.body, tcx); } @@ -9305,7 +9309,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { let elements: [[Option<&ast::Expr>; 1]; 0] = []; let mut infer_element_ty = |_: &mut Self, _| Type::unknown(); - self.infer_collection_literal( + self.try_infer_collection_literal( collection_class, Some(call_expression.into()), &elements, @@ -13038,19 +13042,10 @@ fn is_collection_literal(expression: &ast::Expr) -> bool { ) } -// An empty literal cannot provide element types for the other branch. Infer the -// nonempty branch first so the empty branch can use those types as context. -fn is_nonempty_collection_literal_with_empty_peer( - expression: &ast::Expr, - peer: &ast::Expr, -) -> bool { - match (expression, peer) { - (ast::Expr::List(expression), ast::Expr::List(peer)) => { - !expression.elts.is_empty() && peer.elts.is_empty() - } - (ast::Expr::Dict(expression), ast::Expr::Dict(peer)) => { - !expression.items.is_empty() && peer.items.is_empty() - } +fn is_empty_collection_literal(expression: &ast::Expr) -> bool { + match expression { + ast::Expr::List(list) => list.elts.is_empty(), + ast::Expr::Dict(dict) => dict.items.is_empty(), _ => false, } } diff --git a/crates/ty_python_semantic/src/types/infer/builder/dict.rs b/crates/ty_python_semantic/src/types/infer/builder/dict.rs index cf0e250f78..6c6fa5ba6f 100644 --- a/crates/ty_python_semantic/src/types/infer/builder/dict.rs +++ b/crates/ty_python_semantic/src/types/infer/builder/dict.rs @@ -79,7 +79,7 @@ impl<'db> TypeInferenceBuilder<'db, '_> { // Lower `dict(a=..., b=...)` to synthetic `(Literal["a"], value)` pairs so we can // reuse dict-literal inference. We key the synthetic name off the value node because - // `infer_collection_literal` operates on expressions rather than keywords. + // `try_infer_collection_literal` operates on expressions rather than keywords. let items = arguments .keywords .iter() @@ -106,7 +106,7 @@ impl<'db> TypeInferenceBuilder<'db, '_> { } }; - self.infer_collection_literal( + self.try_infer_collection_literal( KnownClass::Dict, collection_expr, &items,