Add better abstractions to place.rs

Author: AlexWaygood

crates/ty_python_semantic/src/semantic_index/builder.rs

@@ -33,7 +33,7 @@ use crate::semantic_index::definition::{
 use crate::semantic_index::expression::{Expression, ExpressionKind};
 use crate::semantic_index::place::{
     FileScopeId, NodeWithScopeKey, NodeWithScopeKind, NodeWithScopeRef, PlaceExpr,
-    PlaceExprSubSegment, PlaceTableBuilder, Scope, ScopeId, ScopeKind, ScopedPlaceId,
+    PlaceTableBuilder, Scope, ScopeId, ScopeKind, ScopedPlaceId,
 };
 use crate::semantic_index::predicate::{
     PatternPredicate, PatternPredicateKind, Predicate, PredicateNode, ScopedPredicateId,
@@ -1965,7 +1965,7 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
             | ast::Expr::Subscript(ast::ExprSubscript { ctx, .. }) => {
                 if let Ok(mut place_expr) = PlaceExpr::try_from(expr) {
                     if self.is_method_of_class().is_some()
-                        && matches!(place_expr.sub_segments(), &[PlaceExprSubSegment::Member(_)])
+                        && place_expr.is_instance_attribute_candidate()
                     {
                         // We specifically mark attribute assignments to the first parameter of a method,
                         // i.e. typically `self` or `cls`.

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 occured 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 occured 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,12 +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()
-                        .expect("Non-members should never be registered as instance attributes");
-                    self.members.insert(name.clone());
-                }
+                self.members
+                    .extend(place_table.instance_attributes().cloned());
             }
         }
     }