[ty] Don't suggest things that aren't subclasses of BaseException after raise

This only applies to items that have a type associated with them. That
is, things that are already in scope. For items that don't have a type
associated with them (i.e., suggestions from auto-import), we still
suggest them since we can't know if they're appropriate or not. It's not
quite clear on how best to improve here for the auto-import case. (Short
of, say, asking for the type of each such symbol. But the performance
implications of that aren't known yet.)

Note that because of auto-import, we were still suggesting
`NotImplemented` even though astral-sh/ty#1262 specifically cites it as
the motivating example that we *shouldn't* suggest. This was occuring
because auto-import was including symbols from the `builtins` module,
even though those are actually already in scope. So this PR also gets
rid of those suggestions from auto-import.

Overall, this means that, at least, `raise NotImpl` won't suggest
`NotImplemented`.

Fixes astral-sh/ty#1262

Author: BurntSushi

crates/ty_completion_eval/completion-evaluation-tasks.csv

@@ -18,11 +18,11 @@ numpy-array,main.py,1,1
 object-attr-instance-methods,main.py,0,1
 object-attr-instance-methods,main.py,1,1
 pass-keyword-completion,main.py,0,1
-raise-uses-base-exception,main.py,0,2
+raise-uses-base-exception,main.py,0,1
 scope-existing-over-new-import,main.py,0,1
 scope-prioritize-closer,main.py,0,2
 scope-simple-long-identifier,main.py,0,1
 tstring-completions,main.py,0,1
 ty-extensions-lower-stdlib,main.py,0,8
 type-var-typing-over-ast,main.py,0,3
-type-var-typing-over-ast,main.py,1,279
+type-var-typing-over-ast,main.py,1,278

crates/ty_ide/src/completion.rs

@@ -10,8 +10,8 @@ use ruff_python_codegen::Stylist;
 use ruff_python_parser::{Token, TokenAt, TokenKind, Tokens};
 use ruff_text_size::{Ranged, TextLen, TextRange, TextSize};
 use ty_python_semantic::{
-    Completion as SemanticCompletion, ModuleName, NameKind, SemanticModel,
-    types::{CycleDetector, Type},
+    Completion as SemanticCompletion, KnownModule, ModuleName, NameKind, SemanticModel,
+    types::{CycleDetector, KnownClass, Type},
 };
 
 use crate::docstring::Docstring;
@@ -82,6 +82,11 @@ impl<'db> Completions<'db> {
     fn force_add(&mut self, completion: Completion<'db>) {
         self.items.push(completion);
     }
+
+    /// Removes any completion that doesn't satisfy the given predicate.
+    fn retain(&mut self, predicate: impl FnMut(&Completion<'_>) -> bool) {
+        self.items.retain(predicate);
+    }
 }
 
 impl<'db> Extend<SemanticCompletion<'db>> for Completions<'db> {
@@ -364,6 +369,14 @@ pub fn completion<'db>(
         }
     }
 
+    if is_raising_exception(tokens) {
+        let type_base_exception = KnownClass::BaseException.to_subclass_of(db);
+        completions.retain(|c| {
+            let Some(ty) = c.ty else { return true };
+            ty.is_assignable_to(db, type_base_exception)
+        });
+    }
+
     completions.into_completions()
 }
 
@@ -427,7 +440,8 @@ fn add_unimported_completions<'db>(
     let members = importer.members_in_scope_at(scoped.node, scoped.node.start());
 
     for symbol in all_symbols(db, &completions.query) {
-        if symbol.module.file(db) == Some(file) {
+        if symbol.module.file(db) == Some(file) || symbol.module.is_known(db, KnownModule::Builtins)
+        {
             continue;
         }
 
@@ -1358,6 +1372,30 @@ fn is_in_variable_binding(parsed: &ParsedModuleRef, offset: TextSize, typed: Opt
     })
 }
 
+/// Returns true when the cursor is after a `raise` keyword.
+fn is_raising_exception(tokens: &[Token]) -> bool {
+    /// The maximum number of tokens we're willing to
+    /// look-behind to find a `raise` keyword.
+    const LIMIT: usize = 10;
+
+    // This only looks for things like `raise foo.bar.baz.qu<CURSOR>`.
+    // Technically, any kind of expression is allowed after `raise`.
+    // But we may not always wanted to treat it specially. So we're
+    // rather conservative about what we consider "raising an
+    // exception" to be for the purposes of completions. The failure
+    // mode here is that we may wind up suggesting things that
+    // shouldn't be raised. The benefit is that when this heuristic
+    // does work, we won't suggest things that shouldn't be raised.
+    for token in tokens.iter().rev().take(LIMIT) {
+        match token.kind() {
+            TokenKind::Name | TokenKind::Dot => continue,
+            TokenKind::Raise => return true,
+            _ => return false,
+        }
+    }
+    false
+}
+
 /// Order completions according to the following rules:
 ///
 /// 1) Names with no underscore prefix

crates/ty_python_semantic/src/lib.rs

@@ -12,8 +12,8 @@ pub use db::Db;
 pub use diagnostic::add_inferred_python_version_hint_to_diagnostic;
 pub use module_name::{ModuleName, ModuleNameResolutionError};
 pub use module_resolver::{
-    Module, SearchPath, SearchPathValidationError, SearchPaths, all_modules, list_modules,
-    resolve_module, resolve_real_module, system_module_search_paths,
+    KnownModule, Module, SearchPath, SearchPathValidationError, SearchPaths, all_modules,
+    list_modules, resolve_module, resolve_real_module, system_module_search_paths,
 };
 pub use program::{
     Program, ProgramSettings, PythonVersionFileSource, PythonVersionSource,

crates/ty_python_semantic/src/module_resolver/mod.rs

@@ -1,7 +1,7 @@
 use std::iter::FusedIterator;
 
 pub use list::{all_modules, list_modules};
-pub(crate) use module::KnownModule;
+pub use module::KnownModule;
 pub use module::Module;
 pub use path::{SearchPath, SearchPathValidationError};
 pub use resolver::SearchPaths;

crates/ty_python_semantic/src/module_resolver/module.rs

@@ -67,7 +67,7 @@ impl<'db> Module<'db> {
     }
 
     /// Does this module represent the given known module?
-    pub(crate) fn is_known(self, db: &'db dyn Database, known_module: KnownModule) -> bool {
+    pub fn is_known(self, db: &'db dyn Database, known_module: KnownModule) -> bool {
         self.known(db) == Some(known_module)
     }
 

crates/ty_python_semantic/src/types.rs

@@ -74,7 +74,8 @@ use crate::types::variance::VarianceInferable;
 use crate::types::visitor::{any_over_type, exceeds_max_specialization_depth};
 use crate::unpack::EvaluationMode;
 use crate::{Db, FxOrderSet, Module, Program};
-pub(crate) use class::{ClassLiteral, ClassType, GenericAlias, KnownClass};
+pub use class::KnownClass;
+pub(crate) use class::{ClassLiteral, ClassType, GenericAlias};
 use instance::Protocol;
 pub use instance::{NominalInstanceType, ProtocolInstanceType};
 pub use special_form::SpecialFormType;
@@ -1662,8 +1663,8 @@ impl<'db> Type<'db> {
 
     /// Return true if this type is assignable to type `target`.
     ///
-    /// See [`TypeRelation::Assignability`] for more details.
-    pub(crate) fn is_assignable_to(self, db: &'db dyn Db, target: Type<'db>) -> bool {
+    /// See `TypeRelation::Assignability` for more details.
+    pub fn is_assignable_to(self, db: &'db dyn Db, target: Type<'db>) -> bool {
         self.when_assignable_to(db, target, InferableTypeVars::None)
             .is_always_satisfied(db)
     }

crates/ty_python_semantic/src/types/class.rs

@@ -4823,7 +4823,7 @@ impl KnownClass {
     /// representing that class and all possible subclasses of the class.
     ///
     /// If the class cannot be found in typeshed, a debug-level log message will be emitted stating this.
-    pub(crate) fn to_subclass_of(self, db: &dyn Db) -> Type<'_> {
+    pub fn to_subclass_of(self, db: &dyn Db) -> Type<'_> {
         self.to_class_literal(db)
             .to_class_type(db)
             .map(|class| SubclassOfType::from(db, class))