diff --git a/crates/ty_python_semantic/resources/mdtest/assignment/augmented.md b/crates/ty_python_semantic/resources/mdtest/assignment/augmented.md index a572aa0bccf20..160ea0b12aa08 100644 --- a/crates/ty_python_semantic/resources/mdtest/assignment/augmented.md +++ b/crates/ty_python_semantic/resources/mdtest/assignment/augmented.md @@ -16,6 +16,15 @@ x += (3, 4) reveal_type(x) # revealed: tuple[Literal[1, 2, 3, 4], ...] ``` +## Walrus target + +```py +def f(xs: list[int | str]) -> None: + ys = xs + ys[0] = "s" + (ys := [1])[0] += 1 +``` + ## Dunder methods ```py diff --git a/crates/ty_python_semantic/resources/mdtest/attributes.md b/crates/ty_python_semantic/resources/mdtest/attributes.md index 14d812c1b0de2..fdbd23d2b4177 100644 --- a/crates/ty_python_semantic/resources/mdtest/attributes.md +++ b/crates/ty_python_semantic/resources/mdtest/attributes.md @@ -2504,6 +2504,20 @@ class C: C().x ``` +### Walrus reassignment of `self` + +```py +class Other: + x: int = 1 + +class C: + def __init__(self, other: Other) -> None: + (self := other).x = 1 + +# error: [unresolved-attribute] +reveal_type(C(Other()).x) # revealed: Unknown +``` + ### Assignment to `self` after nested function ```py diff --git a/crates/ty_python_semantic/resources/mdtest/expression/attribute.md b/crates/ty_python_semantic/resources/mdtest/expression/attribute.md index 60c1bc518aeea..acd4cfbb7c49a 100644 --- a/crates/ty_python_semantic/resources/mdtest/expression/attribute.md +++ b/crates/ty_python_semantic/resources/mdtest/expression/attribute.md @@ -32,3 +32,18 @@ def _(flag: bool): # error: [unresolved-attribute] "Class `A` has no attribute `non_existent`" reveal_type(A.non_existent) # revealed: Unknown ``` + +## Walrus attribute access after later rebinding + +```py +class IntBox: + attr: int + +class StrBox: + attr: str + +def f() -> None: + (box := IntBox()).attr = 1 + box = StrBox() + reveal_type(box.attr) # revealed: str +``` diff --git a/crates/ty_python_semantic/resources/mdtest/narrow/truthiness.md b/crates/ty_python_semantic/resources/mdtest/narrow/truthiness.md index e19d756472968..806569cd6bbbe 100644 --- a/crates/ty_python_semantic/resources/mdtest/narrow/truthiness.md +++ b/crates/ty_python_semantic/resources/mdtest/narrow/truthiness.md @@ -46,6 +46,27 @@ else: reveal_type(x) # revealed: Literal[b"", b"bar", 0, False, ""] | None | tuple[()] ``` +## Walrus Member Access + +We can narrow on an attribute expression, even when its base is a named expression: + +```py +class Foo: + val: int | None + +if (foo := Foo()).val: + reveal_type(foo.val) # revealed: int & ~AlwaysFalsy +``` + +But we don't pick up stale narrowings from before the assignment in the named expression: + +```py +foo1 = Foo() +foo1.val = None +if (foo1 := Foo()).val: + reveal_type(foo1.val) # revealed: int & ~AlwaysFalsy +``` + ## Function Literals Basically functions are always truthy. diff --git a/crates/ty_python_semantic/resources/mdtest/subscript/lists.md b/crates/ty_python_semantic/resources/mdtest/subscript/lists.md index eeac45cdad6fb..1a0c77cd1369d 100644 --- a/crates/ty_python_semantic/resources/mdtest/subscript/lists.md +++ b/crates/ty_python_semantic/resources/mdtest/subscript/lists.md @@ -33,3 +33,30 @@ x["a" if (y := 2) else 1] = 6 # error: [invalid-assignment] x["a" if (y := 2) else "b"] = 6 ``` + +## Walrus subscript access + +```py +xs: list[int | None] = [1] +xs[0] = None + +reveal_type((xs := [1])[0]) # revealed: int | None +``` + +## Walrus subscript access after rebinding + +```py +def f(xs: list[int | str]) -> None: + ys = xs + ys[0] = "s" + reveal_type((ys := [1])[0]) # revealed: int +``` + +## Walrus subscript access after later rebinding + +```py +def f() -> None: + (ys := [1])[0] = 2 + ys = ["s"] + reveal_type(ys[0]) # revealed: str +``` diff --git a/crates/ty_python_semantic/src/semantic_index/builder.rs b/crates/ty_python_semantic/src/semantic_index/builder.rs index f47c72e69a45b..72ef624743743 100644 --- a/crates/ty_python_semantic/src/semantic_index/builder.rs +++ b/crates/ty_python_semantic/src/semantic_index/builder.rs @@ -660,6 +660,104 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> { self.current_place_table_mut().symbol_mut(id).mark_used(); } + fn record_place_use(&mut self, place_id: ScopedPlaceId, expr: &'ast ast::Expr) { + if let ScopedPlaceId::Symbol(symbol_id) = place_id { + self.mark_symbol_used(symbol_id); + } + let use_id = self.current_ast_ids().record_use(expr); + self.current_use_def_map_mut() + .record_use(place_id, use_id, NodeKey::from_node(expr)); + } + + fn record_place_definition(&mut self, place_id: ScopedPlaceId, expr: &'ast ast::Expr) { + match self.current_assignment() { + Some(CurrentAssignment::Assign { node, unpack }) => { + let assignment = self.add_definition( + place_id, + AssignmentDefinitionNodeRef { + unpack, + value: &node.value, + target: expr, + }, + ); + + self.add_dict_key_assignment_definitions(&node.targets, &node.value, assignment); + } + Some(CurrentAssignment::AnnAssign(ann_assign)) => { + self.add_standalone_type_expression(&ann_assign.annotation); + let assignment = self.add_definition( + place_id, + AnnotatedAssignmentDefinitionNodeRef { + node: ann_assign, + annotation: &ann_assign.annotation, + value: ann_assign.value.as_deref(), + target: expr, + }, + ); + + if let Some(value) = ann_assign.value.as_deref() { + self.add_dict_key_assignment_definitions( + [&*ann_assign.target], + value, + assignment, + ); + } + } + Some(CurrentAssignment::AugAssign(aug_assign)) => { + self.add_definition(place_id, aug_assign); + } + Some(CurrentAssignment::For { node, unpack }) => { + self.add_definition( + place_id, + ForStmtDefinitionNodeRef { + unpack, + iterable: &node.iter, + target: expr, + is_async: node.is_async, + }, + ); + } + Some(CurrentAssignment::Named(named)) => { + // TODO(dhruvmanila): If the current scope is a comprehension, then the + // named expression is implicitly nonlocal. This is yet to be + // implemented. + self.add_definition(place_id, named); + } + Some(CurrentAssignment::Comprehension { + unpack, + node, + first, + }) => { + self.add_definition( + place_id, + ComprehensionDefinitionNodeRef { + unpack, + iterable: &node.iter, + target: expr, + first, + is_async: node.is_async, + }, + ); + } + Some(CurrentAssignment::WithItem { + item, + is_async, + unpack, + }) => { + self.add_definition( + place_id, + WithItemDefinitionNodeRef { + unpack, + context_expr: &item.context_expr, + target: expr, + is_async, + }, + ); + } + None => {} + } + } + fn add_entry_for_definition_key(&mut self, key: DefinitionNodeKey) -> &mut Definitions<'db> { self.definitions_by_node.entry(key).or_default() } @@ -2877,27 +2975,37 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> { ast::Expr::Name(ast::ExprName { ctx, .. }) | ast::Expr::Attribute(ast::ExprAttribute { ctx, .. }) | ast::Expr::Subscript(ast::ExprSubscript { ctx, .. }) => { + // Record place effects after walking the expression. For names, this is + // equivalent because `walk_expr` is a no-op; for attribute/subscript places, + // child evaluation can introduce bindings (for example via walrus operators), + // and those bindings need to exist before we register parent/member associations. + let mut deferred_effects = None; if let Some(mut place_expr) = PlaceExpr::try_from_expr(expr) { - if let Some(method_scope_id) = self.is_method_or_eagerly_executed_in_method() { - if let PlaceExpr::Member(member) = &mut place_expr { - if member.is_instance_attribute_candidate() { - // We specifically mark attribute assignments to the first parameter of a method, - // i.e. typically `self` or `cls`. - // However, we must check that the symbol hasn't been shadowed by an intermediate - // scope (e.g., a comprehension variable: `for self in [...]`). - let accessed_object_refers_to_first_parameter = - self.current_first_parameter_name.is_some_and(|first| { - member.symbol_name() == first - && !self.is_symbol_bound_in_intermediate_eager_scopes( - first, - method_scope_id, - ) - }); - - if accessed_object_refers_to_first_parameter { - member.mark_instance_attribute(); - } - } + if let Some(method_scope_id) = self.is_method_or_eagerly_executed_in_method() + && let PlaceExpr::Member(member) = &mut place_expr + && member.is_instance_attribute_candidate() + && let Some(attribute) = expr.as_attribute_expr() + { + // We specifically mark direct attribute assignments to the first + // parameter of a method, i.e. typically `self` or `cls`. + // However, we must check that the symbol hasn't been shadowed by an + // intermediate scope (e.g., a comprehension variable: `for self in [...]`) + // and that the AST base is still the original name rather than a + // rebinding expression such as `(self := other).x`. + let accessed_object_refers_to_first_parameter = + self.current_first_parameter_name.is_some_and(|first| { + attribute + .value + .as_name_expr() + .is_some_and(|name| name.id == first) + && !self.is_symbol_bound_in_intermediate_eager_scopes( + first, + method_scope_id, + ) + }); + + if accessed_object_refers_to_first_parameter { + member.mark_instance_attribute(); } } @@ -2911,110 +3019,26 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> { (ast::ExprContext::Del, _) => (true, true), (ast::ExprContext::Invalid, _) => (false, false), }; - let place_id = self.add_place(place_expr); + deferred_effects = Some((place_expr, is_use, is_definition)); + } + + // Track reachability of attribute expressions to silence `unresolved-attribute` + // diagnostics in unreachable code. + if expr.is_attribute_expr() { + self.current_use_def_map_mut() + .record_node_reachability(node_key); + } + walk_expr(self, expr); + + if let Some((place_expr, is_use, is_definition)) = deferred_effects { + let place_id = self.add_place(place_expr); if is_use { - if let ScopedPlaceId::Symbol(symbol_id) = place_id { - self.mark_symbol_used(symbol_id); - } - let use_id = self.current_ast_ids().record_use(expr); - self.current_use_def_map_mut() - .record_use(place_id, use_id, node_key); + self.record_place_use(place_id, expr); } - if is_definition { - match self.current_assignment() { - Some(CurrentAssignment::Assign { node, unpack }) => { - let assignment = self.add_definition( - place_id, - AssignmentDefinitionNodeRef { - unpack, - value: &node.value, - target: expr, - }, - ); - - self.add_dict_key_assignment_definitions( - &node.targets, - &node.value, - assignment, - ); - } - Some(CurrentAssignment::AnnAssign(ann_assign)) => { - self.add_standalone_type_expression(&ann_assign.annotation); - let assignment = self.add_definition( - place_id, - AnnotatedAssignmentDefinitionNodeRef { - node: ann_assign, - annotation: &ann_assign.annotation, - value: ann_assign.value.as_deref(), - target: expr, - }, - ); - - if let Some(value) = ann_assign.value.as_deref() { - self.add_dict_key_assignment_definitions( - [&*ann_assign.target], - value, - assignment, - ); - } - } - Some(CurrentAssignment::AugAssign(aug_assign)) => { - self.add_definition(place_id, aug_assign); - } - Some(CurrentAssignment::For { node, unpack }) => { - self.add_definition( - place_id, - ForStmtDefinitionNodeRef { - unpack, - iterable: &node.iter, - target: expr, - is_async: node.is_async, - }, - ); - } - Some(CurrentAssignment::Named(named)) => { - // TODO(dhruvmanila): If the current scope is a comprehension, then the - // named expression is implicitly nonlocal. This is yet to be - // implemented. - self.add_definition(place_id, named); - } - Some(CurrentAssignment::Comprehension { - unpack, - node, - first, - }) => { - self.add_definition( - place_id, - ComprehensionDefinitionNodeRef { - unpack, - iterable: &node.iter, - target: expr, - first, - is_async: node.is_async, - }, - ); - } - Some(CurrentAssignment::WithItem { - item, - is_async, - unpack, - }) => { - self.add_definition( - place_id, - WithItemDefinitionNodeRef { - unpack, - context_expr: &item.context_expr, - target: expr, - is_async, - }, - ); - } - None => {} - } + self.record_place_definition(place_id, expr); } - if let Some(unpack_position) = self .current_assignment_mut() .and_then(CurrentAssignment::unpack_position_mut) @@ -3022,15 +3046,6 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> { *unpack_position = UnpackPosition::Other; } } - - // Track reachability of attribute expressions to silence `unresolved-attribute` - // diagnostics in unreachable code. - if expr.is_attribute_expr() { - self.current_use_def_map_mut() - .record_node_reachability(node_key); - } - - walk_expr(self, expr); } ast::Expr::Named(node) => { // TODO walrus in comprehensions is implicitly nonlocal diff --git a/crates/ty_python_semantic/src/semantic_index/member.rs b/crates/ty_python_semantic/src/semantic_index/member.rs index 05c4f13d67642..f447694fc20a6 100644 --- a/crates/ty_python_semantic/src/semantic_index/member.rs +++ b/crates/ty_python_semantic/src/semantic_index/member.rs @@ -26,13 +26,6 @@ impl Member { } } - /// Returns the left most part of the member expression, e.g. `x` in `x.y.z`. - /// - /// This is the symbol on which the member access is performed. - pub(crate) fn symbol_name(&self) -> &str { - self.expression.symbol_name() - } - pub(crate) fn expression(&self) -> &MemberExpr { &self.expression } @@ -226,6 +219,9 @@ impl MemberExprBuilder { path: name.id.clone(), segments: smallvec::SmallVec::new_const(), }), + ast::ExprRef::Named(named) => { + MemberExprBuilder::visit_expr(ast::ExprRef::from(named.target.as_ref())) + } ast::ExprRef::Attribute(attribute) => { let mut builder = diff --git a/crates/ty_python_semantic/src/types/infer/builder.rs b/crates/ty_python_semantic/src/types/infer/builder.rs index 355336338b388..941bf7759d884 100644 --- a/crates/ty_python_semantic/src/types/infer/builder.rs +++ b/crates/ty_python_semantic/src/types/infer/builder.rs @@ -7646,6 +7646,19 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { return (Place::Undefined, None); } + // A named expression can show up here when resolving the parent place of something + // like `(foo := bar()).baz`. It binds `foo`, but it is not a normal load site and + // therefore has no `ScopedUseId`, so resolve it from its binding definition instead. + if let ast::ExprRef::Named(named) = expr_ref { + let place = if named.target.is_name_expr() { + let definition = self.index.expect_single_definition(named); + Place::bound(binding_type(db, definition)) + } else { + Place::Undefined + }; + return (place, None); + } + let use_id = expr_ref.scoped_use_id(db, scope); let place = place_from_bindings(db, use_def.bindings_at_use(use_id)).place;