diff --git a/crates/ty_ide/src/goto_definition.rs b/crates/ty_ide/src/goto_definition.rs index fc0fc28fb9..81caea650f 100644 --- a/crates/ty_ide/src/goto_definition.rs +++ b/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 = 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 = 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 . /// We must ensure we respect re-import convention for stub files for /// imports in builtins.pyi. diff --git a/crates/ty_ide/src/hover.rs b/crates/ty_ide/src/hover.rs index 1b348b82b9..3b9bf7eeb4 100644 --- a/crates/ty_ide/src/hover.rs +++ b/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 = 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; diff --git a/crates/ty_ide/src/semantic_tokens.rs b/crates/ty_ide/src/semantic_tokens.rs index df1bb88b37..4fb3aa45ab 100644 --- a/crates/ty_ide/src/semantic_tokens.rs +++ b/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( diff --git a/crates/ty_python_semantic/src/types.rs b/crates/ty_python_semantic/src/types.rs index 7f950f7b77..e91007047d 100644 --- a/crates/ty_python_semantic/src/types.rs +++ b/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") diff --git a/crates/ty_python_semantic/src/types/ide_support.rs b/crates/ty_python_semantic/src/types/ide_support.rs index 4b42458c4d..89dfe8fbe8 100644 --- a/crates/ty_python_semantic/src/types/ide_support.rs +++ b/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