[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

@@ -9,9 +9,10 @@ use ruff_python_ast::name::Name;
 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::types::UnionType;
 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 +83,31 @@ impl<'db> Completions<'db> {
     fn force_add(&mut self, completion: Completion<'db>) {
         self.items.push(completion);
     }
+
+    /// Tags completions with whether they are known to be usable in
+    /// a `raise` context.
+    ///
+    /// It's possible that some completions are usable in a `raise`
+    /// but aren't marked by this method. That is, false negatives are
+    /// possible but false positives are not.
+    fn tag_raisable(&mut self) {
+        let raisable_type = UnionType::from_elements(
+            self.db,
+            [
+                KnownClass::BaseException.to_subclass_of(self.db),
+                KnownClass::BaseException.to_instance(self.db),
+            ],
+        );
+        for item in &mut self.items {
+            let Some(ty) = item.ty else { continue };
+            item.is_definitively_raisable = ty.is_assignable_to(self.db, raisable_type);
+        }
+    }
+
+    /// 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> {
@@ -153,6 +179,13 @@ pub struct Completion<'db> {
     /// Whether this item only exists for type checking purposes and
     /// will be missing at runtime
     pub is_type_check_only: bool,
+    /// Whether this item can definitively be used in a `raise` context.
+    ///
+    /// Note that this may not always be computed. (i.e., Only computed
+    /// when we are in a `raise` context.) And also note that if this
+    /// is `true`, then it's definitively usable in `raise`, but if
+    /// it's `false`, it _may_ still be usable in `raise`.
+    pub is_definitively_raisable: bool,
     /// The documentation associated with this item, if
     /// available.
     pub documentation: Option<Docstring>,
@@ -177,6 +210,7 @@ impl<'db> Completion<'db> {
             import: None,
             builtin: semantic.builtin,
             is_type_check_only,
+            is_definitively_raisable: false,
             documentation,
         }
     }
@@ -257,6 +291,7 @@ impl<'db> Completion<'db> {
             import: None,
             builtin: false,
             is_type_check_only: false,
+            is_definitively_raisable: false,
             documentation: None,
         }
     }
@@ -271,6 +306,7 @@ impl<'db> Completion<'db> {
             import: None,
             builtin: true,
             is_type_check_only: false,
+            is_definitively_raisable: false,
             documentation: None,
         }
     }
@@ -364,6 +400,27 @@ pub fn completion<'db>(
         }
     }
 
+    if is_raising_exception(tokens) {
+        completions.tag_raisable();
+
+        // As a special case, and because it's a common footgun, we
+        // specifically disallow `NotImplemented` in this context.
+        // `NotImplementedError` should be used instead. So if we can
+        // definitively detect `NotImplemented`, then we can safely
+        // omit it from suggestions.
+        let not_implemented_type = UnionType::from_elements(
+            db,
+            [
+                KnownClass::NotImplementedType.to_subclass_of(db),
+                KnownClass::NotImplementedType.to_instance(db),
+            ],
+        );
+        completions.retain(|item| {
+            let Some(ty) = item.ty else { return true };
+            !ty.is_assignable_to(db, not_implemented_type)
+        });
+    }
+
     completions.into_completions()
 }
 
@@ -427,7 +484,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;
         }
 
@@ -450,6 +508,7 @@ fn add_unimported_completions<'db>(
             builtin: false,
             // TODO: `is_type_check_only` requires inferring the type of the symbol
             is_type_check_only: false,
+            is_definitively_raisable: false,
             documentation: None,
         });
     }
@@ -1358,6 +1417,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 want 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
@@ -1370,8 +1453,16 @@ fn is_in_variable_binding(parsed: &ParsedModuleRef, offset: TextSize, typed: Opt
 /// This has the effect of putting all dunder attributes after "normal"
 /// attributes, and all single-underscore attributes after dunder attributes.
 fn compare_suggestions(c1: &Completion, c2: &Completion) -> Ordering {
-    fn key<'a>(completion: &'a Completion) -> (bool, bool, bool, NameKind, bool, &'a Name) {
+    fn key<'a>(completion: &'a Completion) -> (bool, bool, bool, bool, NameKind, bool, &'a Name) {
         (
+            // This is only true when we are both in a `raise` context
+            // *and* we know this suggestion is definitively usable
+            // in a `raise` context. So we should sort these before
+            // anything else.
+            !completion.is_definitively_raisable,
+            // When `None`, a completion is for something in the
+            // current module, which we should generally prefer over
+            // something from outside the module.
             completion.module_name.is_some(),
             // At time of writing (2025-11-11), keyword completions
             // are classified as builtins, which makes them sort after

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;
@@ -1670,8 +1671,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)
     }
@@ -12108,7 +12109,7 @@ impl get_size2::GetSize for UnionType<'_> {}
 impl<'db> UnionType<'db> {
     /// Create a union from a list of elements
     /// (which may be eagerly simplified into a different variant of [`Type`] altogether).
-    pub(crate) fn from_elements<I, T>(db: &'db dyn Db, elements: I) -> Type<'db>
+    pub fn from_elements<I, T>(db: &'db dyn Db, elements: I) -> Type<'db>
     where
         I: IntoIterator<Item = T>,
         T: Into<Type<'db>>,

crates/ty_python_semantic/src/types/class.rs

@@ -4743,7 +4743,7 @@ impl KnownClass {
     ///
     /// If the class cannot be found in typeshed, a debug-level log message will be emitted stating this.
     #[track_caller]
-    pub(crate) fn to_instance(self, db: &dyn Db) -> Type<'_> {
+    pub fn to_instance(self, db: &dyn Db) -> Type<'_> {
         debug_assert_ne!(
             self,
             KnownClass::Tuple,
@@ -4896,7 +4896,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))