Repository: astral-sh/ruff
File: crates/ty_ide/src/completion.rs (Rust)

Code before:
        // part of a decorator node.
        if covering_node
            .ancestors()
            // We bail if we're specifying arguments as we don't
            // want to suppress suggestions there.
            .take_while(|node| {
                !matches!(node, ast::AnyNodeRef::Arguments(_)) && !node.is_statement()
            })
            .any(|node| matches!(node, ast::AnyNodeRef::Decorator(_)))
        {
            return Some(FxHashSet::from_iter(["lambda"]));
        }
        covering_node.ancestors().find_map(|node| {
            self.is_in_for_statement_iterable(node)
                .then(|| FxHashSet::from_iter(["yield", "lambda", "await"]))
                .or_else(|| {
                    self.is_expecting_expression(node).then(|| {
                        FxHashSet::from_iter([
                            "await", "lambda", "yield", "for", "if", "else", "and", "or", "not",
                            "in", "is", "True", "False", "None",
                        ])
                    })
                })
        })
    }

    /// Returns true when only an expression is valid after the cursor
    /// according to the python grammar.
    ///
    /// `node` should be the smallest AST node fully covering the
    /// typed text.
    fn is_expecting_expression(&self, node: ast::AnyNodeRef) -> bool {
        let contains = |expr: &ast::Expr| expr.range().contains_range(self.range);

        match node {
            // All checks here are intended to find cases where
            // the python grammar disallows anything but expressions.

            // if_stmt := 'if' named_expression ':' block elif_stmt
            ast::AnyNodeRef::StmtIf(stmt) => {
                contains(&stmt.test)
                    || stmt
                        .elif_else_clauses
                        .iter()
                        .any(|clause| clause.test.as_ref().is_some_and(contains))
            }
            // while_stmt := 'while' named_expression ':' block [else_block]
            ast::AnyNodeRef::StmtWhile(stmt) => contains(&stmt.test),

            // for_stmt := 'for' star_targets 'in' ~ star_expressions ':' [TYPE_COMMENT] block [else_block]
            ast::AnyNodeRef::StmtFor(stmt) => contains(&stmt.iter),

Review comment (refactor):
This leads to us calculating whether or not we're in a class def twice, here and in `add_argument_completions`. I think this is fine for now but we probably want to do some refactors here and store more information on the context such that we only run the detection once and can re-use that both to add completions and for ranking. Perhaps the `NonImportContext` can store another enum and track if we're in a class-def/raising_exception etc.

Have you already thought about how to best scale this @BurntSushi ?

Code after:
            match token.kind() {
                TokenKind::Name | TokenKind::Dot => continue,
                TokenKind::Raise => return true,
                _ => return false,
            }
        }
        false
    }

    /// Returns true when the curser is within the
    /// arguments node of a class definition.
    ///
    /// E.g. `class Foo(Bar<CURSOR>)`
    fn is_in_class_def(&self) -> bool {
        for node in self.covering_node.ancestors() {
            if let ast::AnyNodeRef::StmtClassDef(class_def) = node {
                return class_def
                    .arguments
                    .as_ref()
                    .is_some_and(|args| args.range.contains_range(self.range));
            }
            if node.is_statement() {
                return false;
            }
        }
        false
    }

    /// Returns a set of keywords that are valid at
    /// the current cursor position.
    ///
    /// Returns None if no context-based exclusions can
    /// be identified. Meaning that all keywords are valid.
    fn valid_keywords(&self) -> Option<FxHashSet<&'static str>> {
        // Check if the cursor is within the naming
        // part of a decorator node.
        if self
            .covering_node
            .ancestors()
            // We bail if we're specifying arguments as we don't
            // want to suppress suggestions there.
            .take_while(|node| {
                !matches!(node, ast::AnyNodeRef::Arguments(_)) && !node.is_statement()
            })
            .any(|node| matches!(node, ast::AnyNodeRef::Decorator(_)))
        {
            return Some(FxHashSet::from_iter(["lambda"]));
        }
        self.covering_node.ancestors().find_map(|node| {
            self.is_in_for_statement_iterable(node)
                .then(|| FxHashSet::from_iter(["yield", "lambda", "await"]))