From 93a9fabb26e04a489d40df6138e322c3e9c698da Mon Sep 17 00:00:00 2001
From: Micha Reiser <micha@reiser.io>
Date: Sat, 19 Jul 2025 18:21:27 +0200
Subject: [PATCH] [ty] Avoid secondary tree traversal to get call expression
 for keyword arguments (#19429)

---
 crates/ty_ide/src/find_node.rs |  6 +++++
 crates/ty_ide/src/goto.rs      | 47 ++++++++++++++++++----------------
 2 files changed, 31 insertions(+), 22 deletions(-)

diff --git a/crates/ty_ide/src/find_node.rs b/crates/ty_ide/src/find_node.rs
index 7e8dec4b1b..26c4b489d1 100644
--- a/crates/ty_ide/src/find_node.rs
+++ b/crates/ty_ide/src/find_node.rs
@@ -112,6 +112,12 @@ impl<'a> CoveringNode<'a> {
         Ok(self)
     }
 
+    /// Returns an iterator over the ancestor nodes, starting from the root
+    /// and ending with the covering node.
+    pub(crate) fn ancestors(&self) -> impl Iterator<Item = AnyNodeRef<'a>> + '_ {
+        self.nodes.iter().copied()
+    }
+
     /// Finds the index of the node that fully covers the range and
     /// fulfills the given predicate.
     ///
diff --git a/crates/ty_ide/src/goto.rs b/crates/ty_ide/src/goto.rs
index 101726363a..c7b235b1c2 100644
--- a/crates/ty_ide/src/goto.rs
+++ b/crates/ty_ide/src/goto.rs
@@ -4,7 +4,7 @@ pub use crate::goto_type_definition::goto_type_definition;
 
 use crate::find_node::covering_node;
 use crate::stub_mapping::StubMapper;
-use ruff_db::parsed::{ParsedModuleRef, parsed_module};
+use ruff_db::parsed::ParsedModuleRef;
 use ruff_python_ast::{self as ast, AnyNodeRef};
 use ruff_python_parser::TokenKind;
 use ruff_text_size::{Ranged, TextRange, TextSize};
@@ -40,7 +40,10 @@ pub(crate) enum GotoTarget<'a> {
     /// test(a = 1)
     ///      ^
     /// ```
-    KeywordArgument(&'a ast::Keyword),
+    KeywordArgument {
+        keyword: &'a ast::Keyword,
+        call_expression: &'a ast::ExprCall,
+    },
 
     /// Go to on the rest parameter of a pattern match
     ///
@@ -117,10 +120,10 @@ impl GotoTarget<'_> {
             GotoTarget::Parameter(parameter) => parameter.inferred_type(model),
             GotoTarget::Alias(alias) => alias.inferred_type(model),
             GotoTarget::ExceptVariable(except) => except.inferred_type(model),
-            GotoTarget::KeywordArgument(argument) => {
+            GotoTarget::KeywordArgument { keyword, .. } => {
                 // TODO: Pyright resolves the declared type of the matching parameter. This seems more accurate
                 // than using the inferred value.
-                argument.value.inferred_type(model)
+                keyword.value.inferred_type(model)
             }
             // TODO: Support identifier targets
             GotoTarget::PatternMatchRest(_)
@@ -201,22 +204,13 @@ impl GotoTarget<'_> {
             }
 
             // Handle keyword arguments in call expressions
-            GotoTarget::KeywordArgument(keyword) => {
-                // Find the call expression that contains this keyword
-                let module = parsed_module(db, file).load(db);
-
-                // Use the keyword's range to find the containing call expression
-                let covering_node = covering_node(module.syntax().into(), keyword.range())
-                    .find_first(|node| matches!(node, AnyNodeRef::ExprCall(_)))
-                    .ok()?;
-
-                if let AnyNodeRef::ExprCall(call_expr) = covering_node.node() {
-                    let definitions =
-                        definitions_for_keyword_argument(db, file, keyword, call_expr);
-                    return definitions_to_navigation_targets(db, stub_mapper, definitions);
-                }
-
-                None
+            GotoTarget::KeywordArgument {
+                keyword,
+                call_expression,
+            } => {
+                let definitions =
+                    definitions_for_keyword_argument(db, file, keyword, call_expression);
+                definitions_to_navigation_targets(db, stub_mapper, definitions)
             }
 
             // TODO: Handle multi-part module names in import statements
@@ -237,7 +231,7 @@ impl Ranged for GotoTarget<'_> {
             GotoTarget::Alias(alias) => alias.name.range,
             GotoTarget::ImportedModule(module) => module.module.as_ref().unwrap().range,
             GotoTarget::ExceptVariable(except) => except.name.as_ref().unwrap().range,
-            GotoTarget::KeywordArgument(keyword) => keyword.arg.as_ref().unwrap().range,
+            GotoTarget::KeywordArgument { keyword, .. } => keyword.arg.as_ref().unwrap().range,
             GotoTarget::PatternMatchRest(rest) => rest.rest.as_ref().unwrap().range,
             GotoTarget::PatternKeywordArgument(keyword) => keyword.attr.range,
             GotoTarget::PatternMatchStarName(star) => star.name.as_ref().unwrap().range,
@@ -335,7 +329,16 @@ pub(crate) fn find_goto_target(
             Some(AnyNodeRef::ExceptHandlerExceptHandler(handler)) => {
                 Some(GotoTarget::ExceptVariable(handler))
             }
-            Some(AnyNodeRef::Keyword(keyword)) => Some(GotoTarget::KeywordArgument(keyword)),
+            Some(AnyNodeRef::Keyword(keyword)) => {
+                // Find the containing call expression from the ancestor chain
+                let call_expression = covering_node
+                    .ancestors()
+                    .find_map(ruff_python_ast::AnyNodeRef::expr_call)?;
+                Some(GotoTarget::KeywordArgument {
+                    keyword,
+                    call_expression,
+                })
+            }
             Some(AnyNodeRef::PatternMatchMapping(mapping)) => {
                 Some(GotoTarget::PatternMatchRest(mapping))
             }