diff --git a/crates/ty_python_semantic/resources/mdtest/bidirectional.md b/crates/ty_python_semantic/resources/mdtest/bidirectional.md index c6d255ba0a096..9fe0294621ed7 100644 --- a/crates/ty_python_semantic/resources/mdtest/bidirectional.md +++ b/crates/ty_python_semantic/resources/mdtest/bidirectional.md @@ -43,11 +43,56 @@ def f[T](x: T, cond: bool) -> T | list[T]: l5: int | list[int] = f(1, True) -a: list[int] = [1, 2, *(3, 4, 5)] -reveal_type(a) # revealed: list[int] +x: list[int] = [1, 2, *(3, 4, 5)] +reveal_type(x) # revealed: list[int] -b: list[list[int]] = [[1], [2], *([3], [4])] -reveal_type(b) # revealed: list[list[int]] +x: list[list[int]] = [[1], [2], *([3], [4])] +reveal_type(x) # revealed: list[list[int]] + +x: list[list[int | str]] = [[1], [2]] * 3 +reveal_type(x) # revealed: list[list[int | str]] + +x: list[list[int | str]] = 3 * ([[1]] + [[2]]) +reveal_type(x) # revealed: list[list[int | str]] + +x: list[int | str] = 3 * ["x" for _ in range(3)] +reveal_type(x) # revealed: list[int | str] + +# Tuple elements are inferred individually, but type context can prevent e.g. `int` widening. +x: tuple[list[Literal[1]]] = (list1(1),) +reveal_type(x) # revealed: tuple[list[Literal[1]]] + +x: tuple[list[Literal[1]], ...] = (list1(1),) * 3 +reveal_type(x) # revealed: tuple[list[Literal[1]], ...] + +x: tuple[list[Literal[1]], ...] = 3 * ((list1(1),) + (list1(1),)) +reveal_type(x) # revealed: tuple[list[Literal[1]], ...] + +x: set[int | str] = {1, 2} | {3, 4} +reveal_type(x) # revealed: set[int | str] + +x: set[int | str] = {42 for _ in range(3)} +reveal_type(x) # revealed: set[int | str] + +x: dict[int | str, int | str] = {1: 2} | {3: 4} +reveal_type(x) # revealed: dict[int | str, int | str] + +x: dict[int | str, int | str] = {str(i): i for i in range(3)} +reveal_type(x) # revealed: dict[int | str, int | str] + +# TODO: We currently eagerly pass type context to collection literals on either side of a binary +# operator. That makes the cases above work, but it's not generally sound. For example, it gives the +# wrong result in this case. +class X: + def __add__(self, _: list[int]) -> list[int | str]: + return [] + +# error: [unsupported-operator] "Operator `+` is not supported between objects of type `X` and `list[int | str]`" +x: list[int | str] = X() + [1] + +# TODO: We also don't yet support generic function calls like this. +# error: [invalid-assignment] "Object of type `list[int]` is not assignable to `list[int | str]`" +x: list[int | str] = list1(42) * 3 ``` `typed_dict.py`: @@ -88,6 +133,8 @@ reveal_type(d4_invalid_dict) # revealed: TD d5_literal: dict[Hashable, Callable[..., object]] = {"x": lambda: 1} d5_dict: dict[Hashable, Callable[..., object]] = dict(x=lambda: 1) +d6_dict: TD = {"x": 1} | {"x": 2} + def return_literal() -> TD: return {"x": 1} diff --git a/crates/ty_python_semantic/src/types/infer/builder/binary_expressions.rs b/crates/ty_python_semantic/src/types/infer/builder/binary_expressions.rs index ee71ed89d4b87..10692f70c6028 100644 --- a/crates/ty_python_semantic/src/types/infer/builder/binary_expressions.rs +++ b/crates/ty_python_semantic/src/types/infer/builder/binary_expressions.rs @@ -40,11 +40,11 @@ impl<'db> TypeInferenceBuilder<'db, '_> { node_index: _, } = binary; - let (left_ty, right_ty) = match self.infer_binary_expression_operand_types(left, *op, right) - { - BinaryExpressionOperandTypes::TypedDictResult(ty) => return ty, - BinaryExpressionOperandTypes::Inferred(left_ty, right_ty) => (left_ty, right_ty), - }; + let (left_ty, right_ty) = + match self.infer_binary_expression_operand_types(left, *op, right, tcx) { + BinaryExpressionOperandTypes::TypedDictResult(ty) => return ty, + BinaryExpressionOperandTypes::Inferred(left_ty, right_ty) => (left_ty, right_ty), + }; self.infer_binary_expression_type(binary.into(), false, left_ty, right_ty, *op) .unwrap_or_else(|| { @@ -108,12 +108,37 @@ impl<'db> TypeInferenceBuilder<'db, '_> { left: &ast::Expr, op: ast::Operator, right: &ast::Expr, + tcx: TypeContext<'db>, ) -> BinaryExpressionOperandTypes<'db> { + // As a special case, pass `tcx` to binary operands that are collection literals/displays. + // Note that it's not correct to pass it to all binary operands, for example: + // ``` + // x: list[str] = ["x"] * 3 + // ``` + // It doesn't make sense to pass the list type context to the `3` expression. It wouldn't + // have any effect in this case, but it could in more complicated cases. + // TODO: When we support passing `tcx` through generic method calls, we can remove this + // special case and handle the relevant dunder method instead. + let operand_tcx = |expr: &ast::Expr| -> TypeContext<'db> { + match expr { + ast::Expr::List(_) + | ast::Expr::Tuple(_) + | ast::Expr::Set(_) + | ast::Expr::Dict(_) + | ast::Expr::ListComp(_) + | ast::Expr::SetComp(_) + | ast::Expr::DictComp(_) => tcx, + // Also pass `tcx` to nested binary expressions. + ast::Expr::BinOp(_) => tcx, + _ => TypeContext::default(), + } + }; + // When a dict literal is `|`'d with a TypedDict, infer the non-literal side first // so we can use bidirectional inference on the literal before calling the synthesized // `__or__`/`__ror__` method on the TypedDict side. if op == ast::Operator::BitOr && matches!(left, ast::Expr::Dict(_)) { - let right_ty = self.infer_expression(right, TypeContext::default()); + let right_ty = self.infer_expression(right, operand_tcx(right)); if let Type::TypedDict(typed_dict) = right_ty && let Some(ty) = self.try_typed_dict_pep_584_dunder( left, @@ -128,12 +153,12 @@ impl<'db> TypeInferenceBuilder<'db, '_> { // If the TypedDict update path rejects the literal, fall back to ordinary inference // even though that means re-inferring the literal without TypedDict context. return BinaryExpressionOperandTypes::Inferred( - self.infer_expression(left, TypeContext::default()), + self.infer_expression(left, operand_tcx(left)), right_ty, ); } - let left_ty = self.infer_expression(left, TypeContext::default()); + let left_ty = self.infer_expression(left, operand_tcx(left)); if op == ast::Operator::BitOr && let Type::TypedDict(typed_dict) = left_ty && matches!(right, ast::Expr::Dict(_)) @@ -149,7 +174,7 @@ impl<'db> TypeInferenceBuilder<'db, '_> { BinaryExpressionOperandTypes::Inferred( left_ty, - self.infer_expression(right, TypeContext::default()), + self.infer_expression(right, operand_tcx(right)), ) }