[ty] Fix goto for float and complex in type annotation positions (#21388)

crates/ty_ide/src/goto_definition.rs

@@ -1592,6 +1592,111 @@ a = Test()
         ");
     }
 
+    #[test]
+    fn float_annotation() {
+        let test = CursorTest::builder()
+            .source(
+                "main.py",
+                "
+a: float<CURSOR> = 3.14
+",
+            )
+            .build();
+
+        assert_snapshot!(test.goto_definition(), @r#"
+        info[goto-definition]: Definition
+           --> stdlib/builtins.pyi:346:7
+            |
+        345 | @disjoint_base
+        346 | class int:
+            |       ^^^
+        347 |     """int([x]) -> integer
+        348 |     int(x, base=10) -> integer
+            |
+        info: Source
+         --> main.py:2:4
+          |
+        2 | a: float = 3.14
+          |    ^^^^^
+          |
+
+        info[goto-definition]: Definition
+           --> stdlib/builtins.pyi:659:7
+            |
+        658 | @disjoint_base
+        659 | class float:
+            |       ^^^^^
+        660 |     """Convert a string or number to a floating-point number, if possible."""
+            |
+        info: Source
+         --> main.py:2:4
+          |
+        2 | a: float = 3.14
+          |    ^^^^^
+          |
+        "#);
+    }
+
+    #[test]
+    fn complex_annotation() {
+        let test = CursorTest::builder()
+            .source(
+                "main.py",
+                "
+a: complex<CURSOR> = 3.14
+",
+            )
+            .build();
+
+        assert_snapshot!(test.goto_definition(), @r#"
+        info[goto-definition]: Definition
+           --> stdlib/builtins.pyi:346:7
+            |
+        345 | @disjoint_base
+        346 | class int:
+            |       ^^^
+        347 |     """int([x]) -> integer
+        348 |     int(x, base=10) -> integer
+            |
+        info: Source
+         --> main.py:2:4
+          |
+        2 | a: complex = 3.14
+          |    ^^^^^^^
+          |
+
+        info[goto-definition]: Definition
+           --> stdlib/builtins.pyi:659:7
+            |
+        658 | @disjoint_base
+        659 | class float:
+            |       ^^^^^
+        660 |     """Convert a string or number to a floating-point number, if possible."""
+            |
+        info: Source
+         --> main.py:2:4
+          |
+        2 | a: complex = 3.14
+          |    ^^^^^^^
+          |
+
+        info[goto-definition]: Definition
+           --> stdlib/builtins.pyi:820:7
+            |
+        819 | @disjoint_base
+        820 | class complex:
+            |       ^^^^^^^
+        821 |     """Create a complex number from a string or numbers.
+            |
+        info: Source
+         --> main.py:2:4
+          |
+        2 | a: complex = 3.14
+          |    ^^^^^^^
+          |
+        "#);
+    }
+
     /// Regression test for <https://github.com/astral-sh/ty/issues/1451>.
     /// We must ensure we respect re-import convention for stub files for
     /// imports in builtins.pyi.

crates/ty_ide/src/hover.rs

@@ -2634,6 +2634,32 @@ def ab(a: int, *, c: int):
         ");
     }
 
+    #[test]
+    fn hover_float_annotation() {
+        let test = cursor_test(
+            r#"
+            a: float<CURSOR> = 3.14
+        "#,
+        );
+
+        assert_snapshot!(test.hover(), @r"
+        int | float
+        ---------------------------------------------
+        ```python
+        int | float
+        ```
+        ---------------------------------------------
+        info[hover]: Hovered content is
+         --> main.py:2:4
+          |
+        2 | a: float = 3.14
+          |    ^^^^^- Cursor offset
+          |    |
+          |    source
+          |
+        ");
+    }
+
     impl CursorTest {
         fn hover(&self) -> String {
             use std::fmt::Write;

crates/ty_ide/src/semantic_tokens.rs

@@ -1178,6 +1178,43 @@ result = check(None)
         "#);
     }
 
+    #[test]
+    fn test_builtin_types() {
+        let test = SemanticTokenTest::new(
+            r#"
+            class Test:
+                a: int
+                b: bool
+                c: str
+                d: float  # TODO: Should be Class
+                e: list[int]
+                f: list[float]  # TODO: Should be Class
+                g: int | float  # TODO: float should be Class
+            "#,
+        );
+
+        assert_snapshot!(test.to_snapshot(&test.highlight_file()), @r#"
+        "Test" @ 7..11: Class [definition]
+        "a" @ 17..18: Variable
+        "int" @ 20..23: Class
+        "b" @ 28..29: Variable
+        "bool" @ 31..35: Class
+        "c" @ 40..41: Variable
+        "str" @ 43..46: Class
+        "d" @ 51..52: Variable
+        "float" @ 54..59: Variable
+        "e" @ 89..90: Variable
+        "list" @ 92..96: Class
+        "int" @ 97..100: Class
+        "f" @ 106..107: Variable
+        "list" @ 109..113: Class
+        "float" @ 114..119: Variable
+        "g" @ 150..151: Variable
+        "int" @ 153..156: Class
+        "float" @ 159..164: Variable
+        "#);
+    }
+
     #[test]
     fn test_semantic_tokens_range() {
         let test = SemanticTokenTest::new(

crates/ty_python_semantic/src/types.rs

@@ -1171,7 +1171,6 @@ impl<'db> Type<'db> {
         }
     }
 
-    #[cfg(test)]
     #[track_caller]
     pub(crate) const fn expect_union(self) -> UnionType<'db> {
         self.as_union().expect("Expected a Type::Union variant")

crates/ty_python_semantic/src/types/ide_support.rs

@@ -10,9 +10,9 @@ use crate::semantic_index::scope::ScopeId;
 use crate::semantic_index::{
     attribute_scopes, global_scope, place_table, semantic_index, use_def_map,
 };
-use crate::types::CallDunderError;
 use crate::types::call::{CallArguments, MatchedArgument};
 use crate::types::signatures::Signature;
+use crate::types::{CallDunderError, UnionType};
 use crate::types::{
     ClassBase, ClassLiteral, DynamicType, KnownClass, KnownInstanceType, Type, TypeContext,
     TypeVarBoundOrConstraints, class::CodeGeneratorKind,
@@ -619,6 +619,29 @@ pub fn definitions_for_name<'db>(
         let Some(builtins_scope) = builtins_module_scope(db) else {
             return Vec::new();
         };
+
+        // Special cases for `float` and `complex` in type annotation positions.
+        // We don't know whether we're in a type annotation position, so we'll just ask `Name`'s type,
+        // which resolves to `int | float` or `int | float | complex` if `float` or `complex` is used in
+        // a type annotation position and `float` or `complex` otherwise.
+        //
+        // https://typing.python.org/en/latest/spec/special-types.html#special-cases-for-float-and-complex
+        if matches!(name_str, "float" | "complex")
+            && let Some(union) = name.inferred_type(&SemanticModel::new(db, file)).as_union()
+            && is_float_or_complex_annotation(db, union, name_str)
+        {
+            return union
+                .elements(db)
+                .iter()
+                .filter_map(|ty| ty.as_nominal_instance())
+                .map(|instance| {
+                    let definition = instance.class_literal(db).definition(db);
+                    let parsed = parsed_module(db, definition.file(db));
+                    ResolvedDefinition::FileWithRange(definition.focus_range(db, &parsed.load(db)))
+                })
+                .collect();
+        }
+
         find_symbol_in_scope(db, builtins_scope, name_str)
             .into_iter()
             .filter(|def| def.is_reexported(db))
@@ -636,6 +659,30 @@ pub fn definitions_for_name<'db>(
     }
 }
 
+fn is_float_or_complex_annotation(db: &dyn Db, ty: UnionType, name: &str) -> bool {
+    let float_or_complex_ty = match name {
+        "float" => UnionType::from_elements(
+            db,
+            [
+                KnownClass::Int.to_instance(db),
+                KnownClass::Float.to_instance(db),
+            ],
+        ),
+        "complex" => UnionType::from_elements(
+            db,
+            [
+                KnownClass::Int.to_instance(db),
+                KnownClass::Float.to_instance(db),
+                KnownClass::Complex.to_instance(db),
+            ],
+        ),
+        _ => return false,
+    }
+    .expect_union();
+
+    ty == float_or_complex_ty
+}
+
 /// Returns all resolved definitions for an attribute expression `x.y`.
 /// This function duplicates much of the functionality in the semantic
 /// analyzer, but it has somewhat different behavior so we've decided