WIP: support goto/hover on string annotations

crates/ty_ide/src/doc_highlights.rs

@@ -15,7 +15,7 @@ pub fn document_highlights(
     let module = parsed.load(db);
 
     // Get the definitions for the symbol at the cursor position
-    let goto_target = find_goto_target(&module, offset)?;
+    let goto_target = find_goto_target(db, file, &module, offset)?;
 
     // Use DocumentHighlights mode which limits search to current file only
     references(db, file, &goto_target, ReferencesMode::DocumentHighlights)

crates/ty_ide/src/goto.rs

@@ -7,9 +7,15 @@ use std::borrow::Cow;
 
 use crate::find_node::covering_node;
 use crate::stub_mapping::StubMapper;
+use ruff_db::Db;
+use ruff_db::files::File;
 use ruff_db::parsed::ParsedModuleRef;
+use ruff_db::source::source_text;
+use ruff_python_ast::ExprRef;
 use ruff_python_ast::{self as ast, AnyNodeRef};
 use ruff_python_parser::TokenKind;
+use ruff_python_parser::Tokens;
+use ruff_python_parser::parse_string_annotation;
 use ruff_text_size::{Ranged, TextRange, TextSize};
 use ty_python_semantic::HasDefinition;
 use ty_python_semantic::ImportAliasResolution;
@@ -139,6 +145,22 @@ pub(crate) enum GotoTarget<'a> {
     /// ```
     TypeParamTypeVarTupleName(&'a ast::TypeParamTypeVarTuple),
 
+    /// Go to some name found in a string annotation
+    ///
+    /// ```py
+    /// def my_func(x: "MyType | int"): ...
+    ///                 ^^^^^^
+    /// ```
+    ///
+    /// For various reasons we can't and shouldn't store the sub-AST node,
+    /// and instead store the covering string literal and the name/range
+    /// we parsed out of it.
+    StringAnnotationExprName {
+        string_literal: &'a ast::ExprStringLiteral,
+        name: String,
+        range: TextRange,
+    },
+
     NonLocal {
         identifier: &'a ast::Identifier,
     },
@@ -258,6 +280,10 @@ impl GotoTarget<'_> {
             GotoTarget::ImportModuleAlias { alias } => alias.inferred_type(model),
             GotoTarget::ExceptVariable(except) => except.inferred_type(model),
             GotoTarget::KeywordArgument { keyword, .. } => keyword.value.inferred_type(model),
+            GotoTarget::StringAnnotationExprName { .. } => {
+                // TODO: make a way to ask the inference engine about a sub-expr of a string annotation
+                return None;
+            }
             // TODO: Support identifier targets
             GotoTarget::PatternMatchRest(_)
             | GotoTarget::PatternKeywordArgument(_)
@@ -298,7 +324,7 @@ impl GotoTarget<'_> {
         match self {
             GotoTarget::Expression(expression) => match expression {
                 ast::ExprRef::Name(name) => Some(DefinitionsOrTargets::Definitions(
-                    definitions_for_name(db, file, name),
+                    definitions_for_name(db, file, name.id.as_str(), ExprRef::Name(name)),
                 )),
                 ast::ExprRef::Attribute(attribute) => Some(DefinitionsOrTargets::Definitions(
                     ty_python_semantic::definitions_for_attribute(db, file, attribute),
@@ -306,6 +332,17 @@ impl GotoTarget<'_> {
                 _ => None,
             },
 
+            GotoTarget::StringAnnotationExprName {
+                string_literal,
+                name,
+                ..
+            } => Some(DefinitionsOrTargets::Definitions(definitions_for_name(
+                db,
+                file,
+                name,
+                ExprRef::StringLiteral(string_literal),
+            ))),
+
             // For already-defined symbols, they are their own definitions
             GotoTarget::FunctionDef(function) => {
                 let model = SemanticModel::new(db, file);
@@ -416,7 +453,6 @@ impl GotoTarget<'_> {
                     None
                 }
             }
-
             _ => None,
         }
     }
@@ -486,11 +522,14 @@ impl GotoTarget<'_> {
             }
             GotoTarget::NonLocal { identifier, .. } => Some(Cow::Borrowed(identifier.as_str())),
             GotoTarget::Globals { identifier, .. } => Some(Cow::Borrowed(identifier.as_str())),
+            GotoTarget::StringAnnotationExprName { name, .. } => Some(Cow::Borrowed(name)),
         }
     }
 
     /// Creates a `GotoTarget` from a `CoveringNode` and an offset within the node
     pub(crate) fn from_covering_node<'a>(
+        db: &dyn Db,
+        file: File,
         covering_node: &crate::find_node::CoveringNode<'a>,
         offset: TextSize,
     ) -> Option<GotoTarget<'a>> {
@@ -641,6 +680,51 @@ impl GotoTarget<'_> {
                 }
             },
 
+            AnyNodeRef::ExprStringLiteral(string_literal) => {
+                // If we encounter a string literal, blindly assume it's a string annotation
+                // like `x: "str | int"` by parsing it as a sub-AST.
+                //
+                // TODO: we *really* should be asking the semantic analysis if this is actually
+                // a string annotation, because this behaviour will effect any random string
+                // literal that looks like a type..!
+                let source = source_text(db, file);
+                let sub_ast = parse_string_annotation(
+                    source.as_str(),
+                    string_literal.as_single_part_string()?,
+                )
+                .ok()?;
+
+                // Now we can resume the search for a GotoTarget in the sub-AST.
+                let sub_target = find_goto_target_impl(
+                    db,
+                    file,
+                    sub_ast.tokens(),
+                    sub_ast.syntax().into(),
+                    offset,
+                )?;
+
+                // Our search should only really be considered a success if we found a GotoTarget
+                // for a bare name, as we only care about those things in a string annotation
+                // [CITATION EXTREMELY NEEDED]
+                let GotoTarget::Expression(ExprRef::Name(name)) = sub_target else {
+                    return None;
+                };
+
+                // We *cannot* return this GotoTarget (or the sub-AST node it contains)
+                // but that's ~fine because that node will make basically everything else
+                // freak out because e.g. it has no defined scope in the typechecker.
+                //
+                // Instead we record the covering string literal node and the name/range we parsed.
+                // The string literal will be used in much the same way as ty_python_semantic's
+                // `DeferredExpressionState::InStringAnnotation`.
+                let range = sub_target.range();
+                Some(GotoTarget::StringAnnotationExprName {
+                    string_literal,
+                    name: name.id.to_string(),
+                    range,
+                })
+            }
+
             node => node.as_expr_ref().map(GotoTarget::Expression),
         }
     }
@@ -672,6 +756,7 @@ impl Ranged for GotoTarget<'_> {
             GotoTarget::TypeParamTypeVarTupleName(tuple) => tuple.name.range,
             GotoTarget::NonLocal { identifier, .. } => identifier.range,
             GotoTarget::Globals { identifier, .. } => identifier.range,
+            GotoTarget::StringAnnotationExprName { range, .. } => *range,
         }
     }
 }
@@ -728,12 +813,23 @@ fn definitions_to_navigation_targets<'db>(
     }
 }
 
-pub(crate) fn find_goto_target(
+pub(crate) fn find_goto_target<'a>(
-    parsed: &ParsedModuleRef,
+    db: &dyn Db,
+    file: File,
+    parsed: &'a ParsedModuleRef,
     offset: TextSize,
-) -> Option<GotoTarget<'_>> {
+) -> Option<GotoTarget<'a>> {
-    let token = parsed
+    find_goto_target_impl(db, file, parsed.tokens(), parsed.syntax().into(), offset)
-        .tokens()
+}
+
+fn find_goto_target_impl<'a>(
+    db: &dyn Db,
+    file: File,
+    tokens: &'a Tokens,
+    root: AnyNodeRef<'a>,
+    offset: TextSize,
+) -> Option<GotoTarget<'a>> {
+    let token = tokens
         .at_offset(offset)
         .max_by_key(|token| match token.kind() {
             TokenKind::Name
@@ -744,11 +840,11 @@ pub(crate) fn find_goto_target(
             _ => 0,
         })?;
 
-    let covering_node = covering_node(parsed.syntax().into(), token.range())
+    let covering_node = covering_node(root, token.range())
         .find_first(|node| node.is_identifier() || node.is_expression())
         .ok()?;
 
-    GotoTarget::from_covering_node(&covering_node, offset)
+    GotoTarget::from_covering_node(db, file, &covering_node, offset)
 }
 
 /// Helper function to resolve a module name and create a navigation target.

crates/ty_ide/src/goto_declaration.rs

@@ -16,7 +16,7 @@ pub fn goto_declaration(
     offset: TextSize,
 ) -> Option<RangedValue<NavigationTargets>> {
     let module = parsed_module(db, file).load(db);
-    let goto_target = find_goto_target(&module, offset)?;
+    let goto_target = find_goto_target(db, file, &module, offset)?;
 
     let declaration_targets = goto_target
         .get_definition_targets(file, db, ImportAliasResolution::ResolveAliases)?

crates/ty_ide/src/goto_definition.rs

@@ -17,7 +17,7 @@ pub fn goto_definition(
     offset: TextSize,
 ) -> Option<RangedValue<NavigationTargets>> {
     let module = parsed_module(db, file).load(db);
-    let goto_target = find_goto_target(&module, offset)?;
+    let goto_target = find_goto_target(db, file, &module, offset)?;
 
     let definition_targets = goto_target
         .get_definition_targets(file, db, ImportAliasResolution::ResolveAliases)?

crates/ty_ide/src/goto_references.rs

@@ -16,7 +16,7 @@ pub fn goto_references(
     let module = parsed.load(db);
 
     // Get the definitions for the symbol at the cursor position
-    let goto_target = find_goto_target(&module, offset)?;
+    let goto_target = find_goto_target(db, file, &module, offset)?;
 
     let mode = if include_declaration {
         ReferencesMode::References

crates/ty_ide/src/goto_type_definition.rs

@@ -11,7 +11,7 @@ pub fn goto_type_definition(
     offset: TextSize,
 ) -> Option<RangedValue<NavigationTargets>> {
     let module = parsed_module(db, file).load(db);
-    let goto_target = find_goto_target(&module, offset)?;
+    let goto_target = find_goto_target(db, file, &module, offset)?;
 
     let model = SemanticModel::new(db, file);
     let ty = goto_target.inferred_type(&model)?;
@@ -225,23 +225,7 @@ mod tests {
             "#,
         );
 
-        assert_snapshot!(test.goto_type_definition(), @r#"
+        assert_snapshot!(test.goto_type_definition(), @"No goto target found");
-        info[goto-type-definition]: Type definition
-           --> stdlib/builtins.pyi:911:7
-            |
-        910 | @disjoint_base
-        911 | class str(Sequence[str]):
-            |       ^^^
-        912 |     """str(object='') -> str
-        913 |     str(bytes_or_buffer[, encoding[, errors]]) -> str
-            |
-        info: Source
-         --> main.py:2:10
-          |
-        2 | a: str = "test"
-          |          ^^^^^^
-          |
-        "#);
     }
 
     #[test]

crates/ty_ide/src/hover.rs

@@ -11,7 +11,7 @@ use ty_python_semantic::{DisplaySettings, SemanticModel};
 
 pub fn hover(db: &dyn Db, file: File, offset: TextSize) -> Option<RangedValue<Hover<'_>>> {
     let parsed = parsed_module(db, file).load(db);
-    let goto_target = find_goto_target(&parsed, offset)?;
+    let goto_target = find_goto_target(db, file, &parsed, offset)?;
 
     if let GotoTarget::Expression(expr) = goto_target {
         if expr.is_literal_expr() {

crates/ty_ide/src/references.rs

@@ -282,7 +282,9 @@ impl LocalReferencesFinder<'_> {
         // where the identifier might be a multi-part module name.
         let offset = covering_node.node().start();
 
-        if let Some(goto_target) = GotoTarget::from_covering_node(covering_node, offset) {
+        if let Some(goto_target) =
+            GotoTarget::from_covering_node(self.db, self.file, covering_node, offset)
+        {
             // Get the definitions for this goto target
             if let Some(current_definitions_nav) = goto_target
                 .get_definition_targets(self.file, self.db, ImportAliasResolution::PreserveAliases)

crates/ty_ide/src/rename.rs

@@ -11,7 +11,7 @@ pub fn can_rename(db: &dyn Db, file: File, offset: TextSize) -> Option<ruff_text
     let module = parsed.load(db);
 
     // Get the definitions for the symbol at the offset
-    let goto_target = find_goto_target(&module, offset)?;
+    let goto_target = find_goto_target(db, file, &module, offset)?;
 
     // Don't allow renaming of import module components
     if matches!(
@@ -60,7 +60,7 @@ pub fn rename(
     let module = parsed.load(db);
 
     // Get the definitions for the symbol at the offset
-    let goto_target = find_goto_target(&module, offset)?;
+    let goto_target = find_goto_target(db, file, &module, offset)?;
 
     // Clients shouldn't call us with an empty new name, but just in case...
     if new_name.is_empty() {

crates/ty_python_semantic/src/types/ide_support.rs

@@ -429,13 +429,13 @@ pub fn definition_kind_for_name<'db>(
 pub fn definitions_for_name<'db>(
     db: &'db dyn Db,
     file: File,
-    name: &ast::ExprName,
+    name_str: &str,
+    enclosing_node: ast::ExprRef,
 ) -> Vec<ResolvedDefinition<'db>> {
     let index = semantic_index(db, file);
-    let name_str = name.id.as_str();
 
     // Get the scope for this name expression
-    let file_scope = index.expression_scope_id(&ast::ExprRef::from(name));
+    let file_scope = index.expression_scope_id(&enclosing_node);
 
     let mut all_definitions = Vec::new();