[ty] Fix panic when attempting to provide autocompletions for an instance of a class that assigns attributes to self[0] (#18707)

Author: AlexWaygood

crates/ty_ide/src/completion.rs

@@ -1835,6 +1835,21 @@ f"{f.<CURSOR>
         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.<CURSOR>
+"#,
+        );
+        test.assert_completions_include("orthogonal_direction");
+    }
+
     impl CursorTest {
         fn completions(&self) -> String {
             self.completions_if(|_| true)

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();

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 `<NAME>.<MEMBER>`
+    /// (meaning it *may* be an instance attribute),
+    /// return `Some(<MEMBER>)`. 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 `<NAME>` 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 `<NAME>.<MEMBER>`,
+    /// 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 `<NAME>` 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(<ATTRIBUTE>)` 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<Item = &PlaceExpr> {
+    pub fn instance_attributes(&self) -> impl Iterator<Item = &Name> {
         self.places()
-            .filter(|place_expr| place_expr.is_instance_attribute())
+            .filter_map(|place_expr| place_expr.as_instance_attribute())
     }
 
     /// Returns the place named `name`.

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());
             }
         }
     }