[ty] Determine completion "kind" primarily based on how the symbol was defined

Fixes astral-sh/ty#2102

crates/ty_ide/src/completion.rs

@@ -11,6 +11,7 @@ use ruff_python_ast::{self as ast, AnyNodeRef};
 use ruff_python_codegen::Stylist;
 use ruff_text_size::{Ranged, TextRange, TextSize};
 use rustc_hash::FxHashSet;
+use ty_python_semantic::semantic_index::definition::Definition;
 use ty_python_semantic::types::UnionType;
 use ty_python_semantic::{
     Completion as SemanticCompletion, KnownModule, ModuleName, NameKind, SemanticModel,
@@ -250,6 +251,9 @@ pub struct Completion<'db> {
     /// an unimported symbol. In that case, computing the
     /// type of all such symbols could be quite expensive.
     pub ty: Option<Type<'db>>,
+    /// The first reachable definition for this
+    /// completion, if available.
+    pub first_reachable_definition: Option<Definition<'db>>,
     /// The "kind" of this completion.
     ///
     /// When this is set, it takes priority over any kind
@@ -306,6 +310,7 @@ impl<'db> Completion<'db> {
             qualified: None,
             insert: None,
             ty: semantic.ty,
+            first_reachable_definition: semantic.first_reachable_definition,
             kind: None,
             module_name: None,
             import: None,
@@ -326,7 +331,7 @@ impl<'db> Completion<'db> {
         type CompletionKindVisitor<'db> =
             CycleDetector<CompletionKind, Type<'db>, Option<CompletionKind>>;
 
-        fn imp<'db>(
+        fn by_type<'db>(
             db: &'db dyn Db,
             ty: Type<'db>,
             visitor: &CompletionKindVisitor<'db>,
@@ -361,10 +366,10 @@ impl<'db> Completion<'db> {
                 Type::Union(union) => union
                     .elements(db)
                     .iter()
-                    .find_map(|&ty| imp(db, ty, visitor))?,
+                    .find_map(|&ty| by_type(db, ty, visitor))?,
                 Type::Intersection(intersection) => intersection
                     .iter_positive(db)
-                    .find_map(|ty| imp(db, ty, visitor))?,
+                    .find_map(|ty| by_type(db, ty, visitor))?,
                 Type::Dynamic(_)
                 | Type::Never
                 | Type::SpecialForm(_)
@@ -372,14 +377,45 @@ impl<'db> Completion<'db> {
                 | Type::AlwaysTruthy
                 | Type::AlwaysFalsy => return None,
                 Type::TypeAlias(alias) => {
-                    visitor.visit(ty, || imp(db, alias.value_type(db), visitor))?
+                    visitor.visit(ty, || by_type(db, alias.value_type(db), visitor))?
                 }
             })
         }
-        self.kind.or_else(|| {
-            self.ty
-                .and_then(|ty| imp(db, ty, &CompletionKindVisitor::default()))
-        })
+
+        fn by_definition<'db>(db: &'db dyn Db, def: Definition<'db>) -> Option<CompletionKind> {
+            use ty_python_semantic::semantic_index::definition::DefinitionKind::*;
+            Some(match def.kind(db) {
+                Import(_) | ImportFrom(_) | ImportFromSubmodule(_) => CompletionKind::Module,
+                Function(_) => CompletionKind::Function,
+                Class(_) => CompletionKind::Class,
+                NamedExpression(_)
+                | Assignment(_)
+                | AnnotatedAssignment(_)
+                | AugmentedAssignment(_)
+                | For(_)
+                | Comprehension(_)
+                | VariadicPositionalParameter(_)
+                | VariadicKeywordParameter(_)
+                | Parameter(_)
+                | WithItem(_)
+                | MatchPattern(_)
+                | ExceptHandler(_)
+                | ParamSpec(_) => CompletionKind::Variable,
+                TypeVar(_) | TypeVarTuple(_) => CompletionKind::TypeParameter,
+                // Give up and defer to type.
+                StarImport(_) | TypeAlias(_) => return None,
+            })
+        }
+
+        self.kind
+            .or_else(|| {
+                self.first_reachable_definition
+                    .and_then(|def| by_definition(db, def))
+            })
+            .or_else(|| {
+                self.ty
+                    .and_then(|ty| by_type(db, ty, &CompletionKindVisitor::default()))
+            })
     }
 
     fn keyword(name: &str) -> Self {
@@ -388,6 +424,7 @@ impl<'db> Completion<'db> {
             qualified: None,
             insert: None,
             ty: None,
+            first_reachable_definition: None,
             kind: Some(CompletionKind::Keyword),
             module_name: None,
             import: None,
@@ -404,6 +441,7 @@ impl<'db> Completion<'db> {
             qualified: None,
             insert: None,
             ty: Some(ty),
+            first_reachable_definition: None,
             kind: Some(CompletionKind::Keyword),
             module_name: None,
             import: None,
@@ -1137,6 +1175,7 @@ fn add_function_arg_completions<'db>(
                 qualified: None,
                 insert,
                 ty: p.ty,
+                first_reachable_definition: None,
                 kind: Some(CompletionKind::Variable),
                 module_name: None,
                 import: None,
@@ -1324,6 +1363,7 @@ fn add_unimported_completions<'db>(
             qualified: Some(ast::name::Name::new(qualified)),
             insert: Some(import_action.symbol_text().into()),
             ty: None,
+            first_reachable_definition: None,
             kind: symbol.kind().to_completion_kind(),
             module_name: Some(module_name),
             import: import_action.import().cloned(),
@@ -7508,6 +7548,52 @@ TypedDi<CURSOR>
         );
     }
 
+    #[test]
+    fn local_variable_kind_string() {
+        let builder = completion_test_builder(
+            "\
+zqzqzq = 'Hello ty'
+zqzq<CURSOR>
+",
+        );
+        assert_snapshot!(
+            builder.kinds().skip_auto_import().build().snapshot(),
+            @"zqzqzq :: Variable",
+        );
+    }
+
+    #[test]
+    fn local_variable_kind_int() {
+        let builder = completion_test_builder(
+            "\
+zqzqzq: int = 3
+zqzq<CURSOR>
+",
+        );
+        assert_snapshot!(
+            builder.kinds().skip_auto_import().build().snapshot(),
+            @"zqzqzq :: Variable",
+        );
+    }
+
+    #[test]
+    fn local_variable_kind_class_and_rebind() {
+        let builder = completion_test_builder(
+            "\
+class zqzqzq: pass
+zqzqzq_var = zqzqzq
+zqzq<CURSOR>
+",
+        );
+        assert_snapshot!(
+            builder.kinds().skip_auto_import().build().snapshot(),
+            @r"
+        zqzqzq :: Class
+        zqzqzq_var :: Variable
+        ",
+        );
+    }
+
     /// A way to create a simple single-file (named `main.py`) completion test
     /// builder.
     ///
@@ -7535,6 +7621,7 @@ TypedDi<CURSOR>
         type_signatures: bool,
         imports: bool,
         module_names: bool,
+        kinds: bool,
         // This doesn't seem like a "very complex" type to me... ---AG
         #[allow(clippy::type_complexity)]
         predicate: Option<Box<dyn Fn(&Completion) -> bool>>,
@@ -7568,6 +7655,7 @@ TypedDi<CURSOR>
                 type_signatures: self.type_signatures,
                 imports: self.imports,
                 module_names: self.module_names,
+                kinds: self.kinds,
             }
         }
 
@@ -7643,6 +7731,13 @@ TypedDi<CURSOR>
             self
         }
 
+        /// When set, the "completion kind" for each symbol is included
+        /// in the snapshot (if available).
+        fn kinds(mut self) -> CompletionTestBuilder {
+            self.kinds = true;
+            self
+        }
+
         /// Apply arbitrary filtering to completions.
         fn filter(
             mut self,
@@ -7674,6 +7769,9 @@ TypedDi<CURSOR>
         /// Whether module names should be included in the snapshot
         /// generated by `CompletionTest::snapshot`.
         module_names: bool,
+        /// Whether "completion kinds" should be included in the snapshot
+        /// generated by `CompletionTest::snapshot`.
+        kinds: bool,
     }
 
     impl<'db> CompletionTest<'db> {
@@ -7717,6 +7815,13 @@ TypedDi<CURSOR>
                             snapshot = format!("{snapshot} :: <no import edit>");
                         }
                     }
+                    if self.kinds {
+                        if let Some(kind) = c.kind(self.db) {
+                            snapshot = format!("{snapshot} :: {kind:?}");
+                        } else {
+                            snapshot = format!("{snapshot} :: <no completion kind>");
+                        }
+                    }
                     snapshot
                 })
                 .collect::<Vec<String>>()
@@ -7767,6 +7872,7 @@ TypedDi<CURSOR>
                 type_signatures: false,
                 imports: false,
                 module_names: false,
+                kinds: false,
                 predicate: None,
             }
         }

crates/ty_python_semantic/src/semantic_model.rs

@@ -118,6 +118,7 @@ impl<'db> SemanticModel<'db> {
                 let ty = Type::module_literal(self.db, self.file, module);
                 Completion {
                     name: Name::new(module.name(self.db).as_str()),
+                    first_reachable_definition: None,
                     ty: Some(ty),
                     builtin,
                 }
@@ -164,9 +165,15 @@ impl<'db> SemanticModel<'db> {
         let builtin = module.is_known(self.db, KnownModule::Builtins);
 
         let mut completions = vec![];
-        for Member { name, ty } in all_members(self.db, ty) {
+        for Member {
+            name,
+            first_reachable_definition,
+            ty,
+        } in all_members(self.db, ty)
+        {
             completions.push(Completion {
                 name,
+                first_reachable_definition,
                 ty: Some(ty),
                 builtin,
             });
@@ -187,6 +194,7 @@ impl<'db> SemanticModel<'db> {
             };
             completions.push(Completion {
                 name: Name::new(base),
+                first_reachable_definition: None,
                 ty: Some(ty),
                 builtin,
             });
@@ -204,6 +212,7 @@ impl<'db> SemanticModel<'db> {
             .into_iter()
             .map(|member| Completion {
                 name: member.name,
+                first_reachable_definition: member.first_reachable_definition,
                 ty: Some(member.ty),
                 builtin: false,
             })
@@ -227,6 +236,7 @@ impl<'db> SemanticModel<'db> {
                 all_reachable_members(self.db, file_scope.to_scope_id(self.db, self.file)).map(
                     |memberdef| Completion {
                         name: memberdef.member.name,
+                        first_reachable_definition: Some(memberdef.first_reachable_definition),
                         ty: Some(memberdef.member.ty),
                         builtin: false,
                     },
@@ -410,6 +420,8 @@ impl NameKind {
 pub struct Completion<'db> {
     /// The label shown to the user for this suggestion.
     pub name: Name,
+    /// The definition from which this symbol was extracted.
+    pub first_reachable_definition: Option<Definition<'db>>,
     /// The type of this completion, if available.
     ///
     /// Generally speaking, this is always available

crates/ty_python_semantic/src/types/list_members.rs

@@ -46,6 +46,7 @@ pub(crate) fn all_end_of_scope_members<'db>(
             let symbol = table.symbol(symbol_id);
             let member = Member {
                 name: symbol.name().clone(),
+                first_reachable_definition: Some(first_reachable_definition),
                 ty,
             };
             Some(MemberWithDefinition {
@@ -66,6 +67,7 @@ pub(crate) fn all_end_of_scope_members<'db>(
                 let symbol = table.symbol(symbol_id);
                 let member = Member {
                     name: symbol.name().clone(),
+                    first_reachable_definition: Some(first_reachable_definition),
                     ty,
                 };
                 Some(MemberWithDefinition {
@@ -101,6 +103,7 @@ pub(crate) fn all_reachable_members<'db>(
                             .ignore_possibly_undefined()?;
                         let member = Member {
                             name: symbol.name().clone(),
+                            first_reachable_definition: Some(first_reachable_definition),
                             ty,
                         };
                         Some(MemberWithDefinition {
@@ -117,6 +120,7 @@ pub(crate) fn all_reachable_members<'db>(
                         let ty = place_with_definition.place.ignore_possibly_undefined()?;
                         let member = Member {
                             name: symbol.name().clone(),
+                            first_reachable_definition: Some(first_reachable_definition),
                             ty,
                         };
                         Some(MemberWithDefinition {
@@ -322,13 +326,19 @@ impl<'db> AllMembers<'db> {
                 let use_def_map = use_def_map(db, module_scope);
                 let place_table = place_table(db, module_scope);
 
-                for (symbol_id, _) in use_def_map.all_end_of_scope_symbol_declarations() {
+                for (symbol_id, declarations) in use_def_map.all_end_of_scope_symbol_declarations()
+                {
                     let symbol_name = place_table.symbol(symbol_id).name();
                     let Place::Defined(ty, _, _) =
                         imported_symbol(db, file, symbol_name, None).place
                     else {
                         continue;
                     };
+                    let Some(first_reachable_definition) =
+                        place_from_declarations(db, declarations).first_declaration
+                    else {
+                        continue;
+                    };
 
                     // Filter private symbols from stubs if they appear to be internal types
                     let is_stub_file = file.path(db).extension() == Some("pyi");
@@ -365,6 +375,7 @@ impl<'db> AllMembers<'db> {
 
                     self.members.insert(Member {
                         name: symbol_name.clone(),
+                        first_reachable_definition: Some(first_reachable_definition),
                         ty,
                     });
                 }
@@ -374,7 +385,11 @@ impl<'db> AllMembers<'db> {
                         |submodule_name| {
                             let ty = literal.resolve_submodule(db, &submodule_name)?;
                             let name = submodule_name.clone();
-                            Some(Member { name, ty })
+                            Some(Member {
+                                name,
+                                first_reachable_definition: None,
+                                ty,
+                            })
                         },
                     ));
             }
@@ -421,6 +436,7 @@ impl<'db> AllMembers<'db> {
                 };
                 self.members.insert(Member {
                     name: memberdef.member.name,
+                    first_reachable_definition: Some(memberdef.first_reachable_definition),
                     ty,
                 });
             }
@@ -452,6 +468,7 @@ impl<'db> AllMembers<'db> {
                     };
                     self.members.insert(Member {
                         name: Name::new(name),
+                        first_reachable_definition: None,
                         ty,
                     });
                 }
@@ -469,6 +486,7 @@ impl<'db> AllMembers<'db> {
                 };
                 self.members.insert(Member {
                     name: memberdef.member.name,
+                    first_reachable_definition: Some(memberdef.first_reachable_definition),
                     ty,
                 });
             }
@@ -496,6 +514,7 @@ impl<'db> AllMembers<'db> {
                     if let Place::Defined(synthetic_member, _, _) = ty.member(db, attr).place {
                         self.members.insert(Member {
                             name: Name::from(*attr),
+                            first_reachable_definition: None,
                             ty: synthetic_member,
                         });
                     }
@@ -529,6 +548,7 @@ pub struct MemberWithDefinition<'db> {
 #[derive(Clone, Debug)]
 pub struct Member<'db> {
     pub name: Name,
+    pub first_reachable_definition: Option<Definition<'db>>,
     pub ty: Type<'db>,
 }