store inferred types for string annotation sub-ASTs

crates/ruff_db/src/parsed.rs

@@ -3,8 +3,10 @@ use std::sync::Arc;
 
 use arc_swap::ArcSwapOption;
 use get_size2::GetSize;
-use ruff_python_ast::{AnyRootNodeRef, ModModule, NodeIndex};
+use ruff_python_ast::{AnyRootNodeRef, ModExpression, ModModule, NodeIndex, StringLiteral};
-use ruff_python_parser::{ParseOptions, Parsed, parse_unchecked};
+use ruff_python_parser::{
+    ParseError, ParseOptions, Parsed, parse_string_annotation, parse_unchecked,
+};
 
 use crate::Db;
 use crate::files::File;
@@ -41,6 +43,18 @@ pub fn parsed_module_impl(db: &dyn Db, file: File) -> Parsed<ModModule> {
         .expect("PySourceType always parses into a module")
 }
 
+pub fn parsed_string_annotation(
+    source: &str,
+    string: &StringLiteral,
+) -> Result<Parsed<ModExpression>, ParseError> {
+    let expr = parse_string_annotation(source, string)?;
+
+    // We need the sub-ast of the string annotation to be indexed
+    indexed::ensure_indexed(&expr);
+
+    Ok(expr)
+}
+
 /// A wrapper around a parsed module.
 ///
 /// This type manages instances of the module AST. A particular instance of the AST
@@ -170,12 +184,21 @@ mod indexed {
         pub parsed: Parsed<ModModule>,
     }
 
+    pub fn ensure_indexed(parsed: &Parsed<ModExpression>) {
+        let mut visitor = Visitor {
+            nodes: Some(Vec::new()),
+            index: 0,
+        };
+
+        AnyNodeRef::from(parsed.syntax()).visit_source_order(&mut visitor);
+    }
+
     impl IndexedModule {
         /// Create a new [`IndexedModule`] from the given AST.
         #[allow(clippy::unnecessary_cast)]
         pub fn new(parsed: Parsed<ModModule>) -> Arc<Self> {
             let mut visitor = Visitor {
-                nodes: Vec::new(),
+                nodes: Some(Vec::new()),
                 index: 0,
             };
 
@@ -186,7 +209,7 @@ mod indexed {
 
             AnyNodeRef::from(inner.parsed.syntax()).visit_source_order(&mut visitor);
 
-            let index: Box<[AnyRootNodeRef<'_>]> = visitor.nodes.into_boxed_slice();
+            let index: Box<[AnyRootNodeRef<'_>]> = visitor.nodes.unwrap().into_boxed_slice();
 
             // SAFETY: We cast from `Box<[AnyRootNodeRef<'_>]>` to `Box<[AnyRootNodeRef<'static>]>`,
             // faking the 'static lifetime to create the self-referential struct. The node references
@@ -215,7 +238,7 @@ mod indexed {
     /// A visitor that collects nodes in source order.
     pub struct Visitor<'a> {
         pub index: u32,
-        pub nodes: Vec<AnyRootNodeRef<'a>>,
+        pub nodes: Option<Vec<AnyRootNodeRef<'a>>>,
     }
 
     impl<'a> Visitor<'a> {
@@ -225,7 +248,9 @@ mod indexed {
             AnyRootNodeRef<'a>: From<&'a T>,
         {
             node.node_index().set(NodeIndex::from(self.index));
-            self.nodes.push(AnyRootNodeRef::from(node));
+            if let Some(nodes) = &mut self.nodes {
+                nodes.push(AnyRootNodeRef::from(node));
+            }
             self.index += 1;
         }
     }

crates/ty/docs/rules.md

@@ -8,7 +8,7 @@
 Default level: <a href="../rules.md#rule-levels" title="This lint has a default level of 'error'."><code>error</code></a> ·
 Added in <a href="https://github.com/astral-sh/ty/releases/tag/0.0.1-alpha.1">0.0.1-alpha.1</a> ·
 <a href="https://github.com/astral-sh/ty/issues?q=sort%3Aupdated-desc%20is%3Aissue%20is%3Aopen%20byte-string-type-annotation" target="_blank">Related issues</a> ·
-<a href="https://github.com/astral-sh/ruff/blob/main/crates%2Fty_python_semantic%2Fsrc%2Ftypes%2Fstring_annotation.rs#L36" target="_blank">View source</a>
+<a href="https://github.com/astral-sh/ruff/blob/main/crates%2Fty_python_semantic%2Fsrc%2Ftypes%2Fstring_annotation.rs#L37" target="_blank">View source</a>
 </small>
 
 
@@ -283,7 +283,7 @@ class A:  # Crash at runtime
 Default level: <a href="../rules.md#rule-levels" title="This lint has a default level of 'error'."><code>error</code></a> ·
 Added in <a href="https://github.com/astral-sh/ty/releases/tag/0.0.1-alpha.1">0.0.1-alpha.1</a> ·
 <a href="https://github.com/astral-sh/ty/issues?q=sort%3Aupdated-desc%20is%3Aissue%20is%3Aopen%20escape-character-in-forward-annotation" target="_blank">Related issues</a> ·
-<a href="https://github.com/astral-sh/ruff/blob/main/crates%2Fty_python_semantic%2Fsrc%2Ftypes%2Fstring_annotation.rs#L120" target="_blank">View source</a>
+<a href="https://github.com/astral-sh/ruff/blob/main/crates%2Fty_python_semantic%2Fsrc%2Ftypes%2Fstring_annotation.rs#L121" target="_blank">View source</a>
 </small>
 
 
@@ -295,7 +295,7 @@ TODO #14889
 Default level: <a href="../rules.md#rule-levels" title="This lint has a default level of 'error'."><code>error</code></a> ·
 Added in <a href="https://github.com/astral-sh/ty/releases/tag/0.0.1-alpha.1">0.0.1-alpha.1</a> ·
 <a href="https://github.com/astral-sh/ty/issues?q=sort%3Aupdated-desc%20is%3Aissue%20is%3Aopen%20fstring-type-annotation" target="_blank">Related issues</a> ·
-<a href="https://github.com/astral-sh/ruff/blob/main/crates%2Fty_python_semantic%2Fsrc%2Ftypes%2Fstring_annotation.rs#L11" target="_blank">View source</a>
+<a href="https://github.com/astral-sh/ruff/blob/main/crates%2Fty_python_semantic%2Fsrc%2Ftypes%2Fstring_annotation.rs#L12" target="_blank">View source</a>
 </small>
 
 
@@ -326,7 +326,7 @@ def test(): -> "int":
 Default level: <a href="../rules.md#rule-levels" title="This lint has a default level of 'error'."><code>error</code></a> ·
 Added in <a href="https://github.com/astral-sh/ty/releases/tag/0.0.1-alpha.1">0.0.1-alpha.1</a> ·
 <a href="https://github.com/astral-sh/ty/issues?q=sort%3Aupdated-desc%20is%3Aissue%20is%3Aopen%20implicit-concatenated-string-type-annotation" target="_blank">Related issues</a> ·
-<a href="https://github.com/astral-sh/ruff/blob/main/crates%2Fty_python_semantic%2Fsrc%2Ftypes%2Fstring_annotation.rs#L86" target="_blank">View source</a>
+<a href="https://github.com/astral-sh/ruff/blob/main/crates%2Fty_python_semantic%2Fsrc%2Ftypes%2Fstring_annotation.rs#L87" target="_blank">View source</a>
 </small>
 
 
@@ -1386,7 +1386,7 @@ super(B, A)  # error: `A` does not satisfy `issubclass(A, B)`
 Default level: <a href="../rules.md#rule-levels" title="This lint has a default level of 'error'."><code>error</code></a> ·
 Added in <a href="https://github.com/astral-sh/ty/releases/tag/0.0.1-alpha.1">0.0.1-alpha.1</a> ·
 <a href="https://github.com/astral-sh/ty/issues?q=sort%3Aupdated-desc%20is%3Aissue%20is%3Aopen%20invalid-syntax-in-forward-annotation" target="_blank">Related issues</a> ·
-<a href="https://github.com/astral-sh/ruff/blob/main/crates%2Fty_python_semantic%2Fsrc%2Ftypes%2Fstring_annotation.rs#L111" target="_blank">View source</a>
+<a href="https://github.com/astral-sh/ruff/blob/main/crates%2Fty_python_semantic%2Fsrc%2Ftypes%2Fstring_annotation.rs#L112" target="_blank">View source</a>
 </small>
 
 
@@ -1859,7 +1859,7 @@ f(x=1)  # Error raised here
 Default level: <a href="../rules.md#rule-levels" title="This lint has a default level of 'error'."><code>error</code></a> ·
 Added in <a href="https://github.com/astral-sh/ty/releases/tag/0.0.1-alpha.1">0.0.1-alpha.1</a> ·
 <a href="https://github.com/astral-sh/ty/issues?q=sort%3Aupdated-desc%20is%3Aissue%20is%3Aopen%20raw-string-type-annotation" target="_blank">Related issues</a> ·
-<a href="https://github.com/astral-sh/ruff/blob/main/crates%2Fty_python_semantic%2Fsrc%2Ftypes%2Fstring_annotation.rs#L61" target="_blank">View source</a>
+<a href="https://github.com/astral-sh/ruff/blob/main/crates%2Fty_python_semantic%2Fsrc%2Ftypes%2Fstring_annotation.rs#L62" target="_blank">View source</a>
 </small>
 
 

crates/ty_ide/src/goto.rs

@@ -313,18 +313,11 @@ impl GotoTarget<'_> {
                 subrange,
                 ..
             } => {
-                let (subast, _submodel) = model.enter_string_annotation(string_expr)?;
+                let (subast, submodel) = model.enter_string_annotation(string_expr)?;
-                let submod = subast.syntax();
+                let subexpr = covering_node(subast.syntax().into(), *subrange)
-                let subnode = covering_node(submod.into(), *subrange).node();
+                    .node()
-
+                    .as_expr_ref()?;
-                // The type checker knows the type of the full annotation but nothing else
+                subexpr.inferred_type(&submodel)
-                if AnyNodeRef::from(&*submod.body) == subnode {
-                    string_expr.inferred_type(model)
-                } else {
-                    // TODO: force the typechecker to tell us its secrets
-                    // (it computes but then immediately discards these types)
-                    return None;
-                }
             }
             GotoTarget::BinOp { expression, .. } => {
                 let (_, ty) = ty_python_semantic::definitions_for_bin_op(model, expression)?;

crates/ty_ide/src/goto_type_definition.rs

@@ -787,7 +787,25 @@ mod tests {
         "#,
         );
 
-        assert_snapshot!(test.goto_type_definition(), @"No goto target found");
+        assert_snapshot!(test.goto_type_definition(), @r#"
+        info[goto-type-definition]: Type definition
+         --> main.py:4:7
+          |
+        2 | a: "None | MyClass" = 1
+        3 |
+        4 | class MyClass:
+          |       ^^^^^^^
+        5 |     """some docs"""
+          |
+        info: Source
+         --> main.py:2:12
+          |
+        2 | a: "None | MyClass" = 1
+          |            ^^^^^^^
+        3 |
+        4 | class MyClass:
+          |
+        "#);
     }
 
     #[test]
@@ -851,7 +869,25 @@ mod tests {
         "#,
         );
 
-        assert_snapshot!(test.goto_type_definition(), @"No goto target found");
+        assert_snapshot!(test.goto_type_definition(), @r#"
+        info[goto-type-definition]: Type definition
+         --> main.py:4:7
+          |
+        2 | a: "None | MyClass" = 1
+        3 |
+        4 | class MyClass:
+          |       ^^^^^^^
+        5 |     """some docs"""
+          |
+        info: Source
+         --> main.py:2:12
+          |
+        2 | a: "None | MyClass" = 1
+          |            ^^^^^^^
+        3 |
+        4 | class MyClass:
+          |
+        "#);
     }
 
     #[test]
@@ -947,7 +983,25 @@ mod tests {
         "#,
         );
 
-        assert_snapshot!(test.goto_type_definition(), @"No goto target found");
+        assert_snapshot!(test.goto_type_definition(), @r#"
+        info[goto-type-definition]: Type definition
+         --> main.py:4:7
+          |
+        2 | a: "MyClass | No" = 1
+        3 |
+        4 | class MyClass:
+          |       ^^^^^^^
+        5 |     """some docs"""
+          |
+        info: Source
+         --> main.py:2:5
+          |
+        2 | a: "MyClass | No" = 1
+          |     ^^^^^^^
+        3 |
+        4 | class MyClass:
+          |
+        "#);
     }
 
     #[test]
@@ -961,7 +1015,25 @@ mod tests {
         "#,
         );
 
-        assert_snapshot!(test.goto_type_definition(), @"No goto target found");
+        assert_snapshot!(test.goto_type_definition(), @r#"
+        info[goto-type-definition]: Type definition
+          --> stdlib/ty_extensions.pyi:20:1
+           |
+        19 | # Types
+        20 | Unknown = object()
+           | ^^^^^^^
+        21 | AlwaysTruthy = object()
+        22 | AlwaysFalsy = object()
+           |
+        info: Source
+         --> main.py:2:15
+          |
+        2 | a: "MyClass | No" = 1
+          |               ^^
+        3 |
+        4 | class MyClass:
+          |
+        "#);
     }
 
     #[test]

crates/ty_ide/src/hover.rs

@@ -953,9 +953,15 @@ mod tests {
         );
 
         assert_snapshot!(test.hover(), @r#"
+        MyClass
+        ---------------------------------------------
         some docs
 
         ---------------------------------------------
+        ```python
+        MyClass
+        ```
+        ---
         some docs
         ---------------------------------------------
         info[hover]: Hovered content is
@@ -998,9 +1004,15 @@ mod tests {
         );
 
         assert_snapshot!(test.hover(), @r#"
+        MyClass
+        ---------------------------------------------
         some docs
 
         ---------------------------------------------
+        ```python
+        MyClass
+        ```
+        ---
         some docs
         ---------------------------------------------
         info[hover]: Hovered content is
@@ -1056,9 +1068,15 @@ mod tests {
         );
 
         assert_snapshot!(test.hover(), @r#"
+        MyClass
+        ---------------------------------------------
         some docs
 
         ---------------------------------------------
+        ```python
+        MyClass
+        ```
+        ---
         some docs
         ---------------------------------------------
         info[hover]: Hovered content is
@@ -1086,7 +1104,25 @@ mod tests {
         "#,
         );
 
-        assert_snapshot!(test.hover(), @"Hover provided no content");
+        assert_snapshot!(test.hover(), @r#"
+        Unknown
+        ---------------------------------------------
+        ```python
+        Unknown
+        ```
+        ---------------------------------------------
+        info[hover]: Hovered content is
+         --> main.py:2:15
+          |
+        2 | a: "MyClass | No" = 1
+          |               ^-
+          |               ||
+          |               |Cursor offset
+          |               source
+        3 |
+        4 | class MyClass:
+          |
+        "#);
     }
 
     #[test]

crates/ty_python_semantic/src/semantic_index/ast_ids.rs

@@ -116,29 +116,44 @@ pub(crate) mod node_key {
     use crate::node_key::NodeKey;
 
     #[derive(Copy, Clone, Eq, PartialEq, Hash, Debug, salsa::Update, get_size2::GetSize)]
-    pub(crate) struct ExpressionNodeKey(NodeKey);
+    pub(crate) struct ExpressionNodeKey(NodeKey, Option<NodeKey>);
+
+    impl From<(NodeKey, &ast::Expr)> for ExpressionNodeKey {
+        fn from(value: (NodeKey, &ast::Expr)) -> Self {
+            Self(value.0, Some(NodeKey::from_node(value.1)))
+        }
+    }
+
+    impl From<(ast::ExprRef<'_>, ast::ExprRef<'_>)> for ExpressionNodeKey {
+        fn from(value: (ast::ExprRef<'_>, ast::ExprRef<'_>)) -> Self {
+            Self(
+                NodeKey::from_node(value.0),
+                Some(NodeKey::from_node(value.1)),
+            )
+        }
+    }
 
     impl From<ast::ExprRef<'_>> for ExpressionNodeKey {
         fn from(value: ast::ExprRef<'_>) -> Self {
-            Self(NodeKey::from_node(value))
+            Self(NodeKey::from_node(value), None)
         }
     }
 
     impl From<&ast::Expr> for ExpressionNodeKey {
         fn from(value: &ast::Expr) -> Self {
-            Self(NodeKey::from_node(value))
+            Self(NodeKey::from_node(value), None)
         }
     }
 
     impl From<&ast::ExprCall> for ExpressionNodeKey {
         fn from(value: &ast::ExprCall) -> Self {
-            Self(NodeKey::from_node(value))
+            Self(NodeKey::from_node(value), None)
         }
     }
 
     impl From<&ast::Identifier> for ExpressionNodeKey {
         fn from(value: &ast::Identifier) -> Self {
-            Self(NodeKey::from_node(value))
+            Self(NodeKey::from_node(value), None)
         }
     }
 }

crates/ty_python_semantic/src/semantic_model.rs

@@ -1,4 +1,5 @@
 use ruff_db::files::{File, FilePath};
+use ruff_db::parsed::parsed_string_annotation;
 use ruff_db::source::{line_index, source_text};
 use ruff_python_ast::{self as ast, ExprStringLiteral, ModExpression};
 use ruff_python_ast::{Expr, ExprRef, HasNodeIndex, name::Name};
@@ -287,6 +288,12 @@ impl<'db> SemanticModel<'db> {
         }
     }
 
+    fn expr_ref_pair<'a>(&'a self, arg: ExprRef<'a>) -> Option<(ExprRef<'a>, ExprRef<'a>)> {
+        self.in_string_annotation_expr
+            .as_ref()
+            .map(|in_ast_expr| (ExprRef::from(in_ast_expr), arg))
+    }
+
     /// Given a string expression, determine if it's a string annotation, and if it is,
     /// yield the parsed sub-AST and a sub-model that knows it's analyzing a sub-AST.
     ///
@@ -317,8 +324,7 @@ impl<'db> SemanticModel<'db> {
         // are not in the File's AST!
         let source = source_text(self.db, self.file);
         let string_literal = string_expr.as_single_part_string()?;
-        let ast =
+        let ast = parsed_string_annotation(source.as_str(), string_literal).ok()?;
-            ruff_python_parser::parse_string_annotation(source.as_str(), string_literal).ok()?;
         let model = Self {
             db: self.db,
             file: self.file,
@@ -422,8 +428,12 @@ impl HasType for ast::ExprRef<'_> {
             return Type::unknown();
         };
         let scope = file_scope.to_scope_id(model.db, model.file);
-
+        let types = infer_scope_types(model.db, scope);
-        infer_scope_types(model.db, scope).expression_type(*self)
+        if let Some((expr_in_ast, sub_expr)) = model.expr_ref_pair(*self) {
+            types.expression_type((expr_in_ast, sub_expr))
+        } else {
+            types.expression_type(model.expr_ref_in_ast(*self))
+        }
     }
 }
 

crates/ty_python_semantic/src/types/infer/builder.rs

@@ -7103,12 +7103,11 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
 
     #[track_caller]
     fn store_expression_type(&mut self, expression: &ast::Expr, ty: Type<'db>) {
-        if self.deferred_state.in_string_annotation() {
+        let expression_key = if let Some(node) = self.deferred_state.active_string_annotation() {
-            // Avoid storing the type of expressions that are part of a string annotation because
+            ExpressionNodeKey::from((node, expression))
-            // the expression ids don't exists in the semantic index. Instead, we'll store the type
+        } else {
-            // on the string expression itself that represents the annotation.
+            ExpressionNodeKey::from(expression)
-            return;
+        };
-        }
 
         let db = self.db();
 
@@ -7116,17 +7115,20 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             MultiInferenceState::Ignore => {}
 
             MultiInferenceState::Panic => {
-                let previous = self.expressions.insert(expression.into(), ty);
+                let previous = self.expressions.insert(expression_key, ty);
-                assert_eq!(previous, None);
+                assert_eq!(
+                    previous, None,
+                    "duplicate key {expression_key:?} for {expression:?}"
+                );
             }
 
             MultiInferenceState::Overwrite => {
-                self.expressions.insert(expression.into(), ty);
+                self.expressions.insert(expression_key, ty);
             }
 
             MultiInferenceState::Intersect => {
                 self.expressions
-                    .entry(expression.into())
+                    .entry(expression_key)
                     .and_modify(|current| {
                         *current = IntersectionType::from_elements(db, [*current, ty]);
                     })
@@ -12306,6 +12308,14 @@ impl DeferredExpressionState {
     const fn in_string_annotation(self) -> bool {
         matches!(self, DeferredExpressionState::InStringAnnotation(_))
     }
+
+    const fn active_string_annotation(self) -> Option<NodeKey> {
+        match self {
+            DeferredExpressionState::None => None,
+            DeferredExpressionState::Deferred => None,
+            DeferredExpressionState::InStringAnnotation(node_key) => Some(node_key),
+        }
+    }
 }
 
 impl From<bool> for DeferredExpressionState {

crates/ty_python_semantic/src/types/string_annotation.rs

@@ -1,3 +1,4 @@
+use ruff_db::parsed::parsed_string_annotation;
 use ruff_db::source::source_text;
 use ruff_python_ast::{self as ast, ModExpression};
 use ruff_python_parser::Parsed;
@@ -149,7 +150,7 @@ pub(crate) fn parse_string_annotation(
         // Compare the raw contents (without quotes) of the expression with the parsed contents
         // contained in the string literal.
         } else if &source[string_literal.content_range()] == string_literal.as_str() {
-            match ruff_python_parser::parse_string_annotation(source.as_str(), string_literal) {
+            match parsed_string_annotation(source.as_str(), string_literal) {
                 Ok(parsed) => return Some(parsed),
                 Err(parse_error) => {
                     if let Some(builder) =