diff --git a/crates/ty_ide/src/completion.rs b/crates/ty_ide/src/completion.rs index a8b93080b4..5fa6cb09f9 100644 --- a/crates/ty_ide/src/completion.rs +++ b/crates/ty_ide/src/completion.rs @@ -1835,6 +1835,21 @@ f"{f. test.assert_completions_include("method"); } + #[test] + fn no_panic_for_attribute_table_that_contains_subscript() { + let test = cursor_test( + r#" +class Point: + def orthogonal_direction(self): + self[0].is_zero + +def test_point(p2: Point): + p2. +"#, + ); + test.assert_completions_include("orthogonal_direction"); + } + impl CursorTest { fn completions(&self) -> String { self.completions_if(|_| true) diff --git a/crates/ty_python_semantic/src/semantic_index/builder.rs b/crates/ty_python_semantic/src/semantic_index/builder.rs index b8e3dc5d0e..c63c357d9f 100644 --- a/crates/ty_python_semantic/src/semantic_index/builder.rs +++ b/crates/ty_python_semantic/src/semantic_index/builder.rs @@ -1964,12 +1964,14 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> { | ast::Expr::Attribute(ast::ExprAttribute { ctx, .. }) | ast::Expr::Subscript(ast::ExprSubscript { ctx, .. }) => { if let Ok(mut place_expr) = PlaceExpr::try_from(expr) { - if self.is_method_of_class().is_some() { + if self.is_method_of_class().is_some() + && place_expr.is_instance_attribute_candidate() + { // We specifically mark attribute assignments to the first parameter of a method, // i.e. typically `self` or `cls`. let accessed_object_refers_to_first_parameter = self .current_first_parameter_name - .is_some_and(|fst| place_expr.root_name().as_str() == fst); + .is_some_and(|fst| place_expr.root_name() == fst); if accessed_object_refers_to_first_parameter && place_expr.is_member() { place_expr.mark_instance_attribute(); diff --git a/crates/ty_python_semantic/src/semantic_index/place.rs b/crates/ty_python_semantic/src/semantic_index/place.rs index f8fa9e1474..edbff98286 100644 --- a/crates/ty_python_semantic/src/semantic_index/place.rs +++ b/crates/ty_python_semantic/src/semantic_index/place.rs @@ -191,16 +191,59 @@ impl PlaceExpr { &self.root_name } + /// If the place expression has the form `.` + /// (meaning it *may* be an instance attribute), + /// return `Some()`. Else, return `None`. + /// + /// This method is internal to the semantic-index submodule. + /// It *only* checks that the AST structure of the `Place` is + /// correct. It does not check whether the `Place` actually occurred in + /// a method context, or whether the `` actually refers to the first + /// parameter of the method (i.e. `self`). To answer those questions, + /// use [`Self::as_instance_attribute`]. + pub(super) fn as_instance_attribute_candidate(&self) -> Option<&Name> { + if self.sub_segments.len() == 1 { + self.sub_segments[0].as_member() + } else { + None + } + } + + /// Return `true` if the place expression has the form `.`, + /// indicating that it *may* be an instance attribute if we are in a method context. + /// + /// This method is internal to the semantic-index submodule. + /// It *only* checks that the AST structure of the `Place` is + /// correct. It does not check whether the `Place` actually occurred in + /// a method context, or whether the `` actually refers to the first + /// parameter of the method (i.e. `self`). To answer those questions, + /// use [`Self::is_instance_attribute`]. + pub(super) fn is_instance_attribute_candidate(&self) -> bool { + self.as_instance_attribute_candidate().is_some() + } + /// Does the place expression have the form `self.{name}` (`self` is the first parameter of the method)? pub(super) fn is_instance_attribute_named(&self, name: &str) -> bool { - self.is_instance_attribute() - && self.sub_segments.len() == 1 - && self.sub_segments[0].as_member().unwrap().as_str() == name + self.as_instance_attribute().map(Name::as_str) == Some(name) } /// Is the place an instance attribute? - pub fn is_instance_attribute(&self) -> bool { - self.flags.contains(PlaceFlags::IS_INSTANCE_ATTRIBUTE) + pub(crate) fn is_instance_attribute(&self) -> bool { + let is_instance_attribute = self.flags.contains(PlaceFlags::IS_INSTANCE_ATTRIBUTE); + if is_instance_attribute { + debug_assert!(self.is_instance_attribute_candidate()); + } + is_instance_attribute + } + + /// Return `Some()` if the place expression is an instance attribute. + pub(crate) fn as_instance_attribute(&self) -> Option<&Name> { + if self.is_instance_attribute() { + debug_assert!(self.as_instance_attribute_candidate().is_some()); + self.as_instance_attribute_candidate() + } else { + None + } } /// Is the place used in its containing scope? @@ -570,9 +613,9 @@ impl PlaceTable { self.places().filter(|place_expr| place_expr.is_name()) } - pub fn instance_attributes(&self) -> impl Iterator { + pub fn instance_attributes(&self) -> impl Iterator { self.places() - .filter(|place_expr| place_expr.is_instance_attribute()) + .filter_map(|place_expr| place_expr.as_instance_attribute()) } /// Returns the place named `name`. diff --git a/crates/ty_python_semantic/src/types/ide_support.rs b/crates/ty_python_semantic/src/types/ide_support.rs index 296c9266d5..1e491e478e 100644 --- a/crates/ty_python_semantic/src/types/ide_support.rs +++ b/crates/ty_python_semantic/src/types/ide_support.rs @@ -185,10 +185,8 @@ impl AllMembers { let index = semantic_index(db, file); for function_scope_id in attribute_scopes(db, class_body_scope) { let place_table = index.place_table(function_scope_id); - for instance_attribute in place_table.instance_attributes() { - let name = instance_attribute.sub_segments()[0].as_member().unwrap(); - self.members.insert(name.clone()); - } + self.members + .extend(place_table.instance_attributes().cloned()); } } }