diff --git a/crates/ty_ide/src/completion.rs b/crates/ty_ide/src/completion.rs index b03b1930ca..ca2305df0c 100644 --- a/crates/ty_ide/src/completion.rs +++ b/crates/ty_ide/src/completion.rs @@ -6433,6 +6433,155 @@ collabc assert_snapshot!(snapshot, @"collections.abc"); } + #[test] + fn local_function_variable_with_return() { + let builder = completion_test_builder( + "\ +variable_global = 1 +def foo(): + variable_local = 1 + variable_ + return +", + ); + assert_snapshot!( + builder.skip_auto_import().build().snapshot(), + @r" + variable_global + variable_local + ", + ); + } + + #[test] + fn nested_scopes_with_return() { + let builder = completion_test_builder( + "\ +variable_1 = 1 +def fun1(): + variable_2 = 1 + def fun2(): + variable_3 = 1 + def fun3(): + variable_4 = 1 + variable_ + return + return + return +", + ); + assert_snapshot!( + builder.skip_auto_import().build().snapshot(), + @r" + variable_1 + variable_2 + variable_3 + variable_4 + ", + ); + } + + #[test] + fn multiple_declarations_global_scope1() { + let builder = completion_test_builder( + "\ +zqzqzq: int = 1 +zqzqzq: str = 'foo' +zqzq +", + ); + // The type for `zqzqzq` *should* be `str`, but we consider all + // reachable declarations and bindings, which means we get a + // union of `int` and `str` here even though the `int` binding + // isn't live at the cursor position. + assert_snapshot!( + builder.skip_auto_import().type_signatures().build().snapshot(), + @"zqzqzq :: int | str", + ); + } + + #[test] + fn multiple_declarations_global_scope2() { + let builder = completion_test_builder( + "\ +zqzqzq: int = 1 +zqzq +zqzqzq: str = 'foo' +", + ); + // The type for `zqzqzq` *should* be `int`, but we consider all + // reachable declarations and bindings, which means we get a + // union of `int` and `str` here even though the `str` binding + // doesn't exist at the cursor position. + assert_snapshot!( + builder.skip_auto_import().type_signatures().build().snapshot(), + @"zqzqzq :: int | str", + ); + } + + #[test] + fn multiple_declarations_function_scope1() { + let builder = completion_test_builder( + "\ +def foo(): + zqzqzq: int = 1 + zqzqzq: str = 'foo' + zqzq + return +", + ); + // The type for `zqzqzq` *should* be `str`, but we consider all + // reachable declarations and bindings, which means we get a + // union of `int` and `str` here even though the `int` binding + // isn't live at the cursor position. + assert_snapshot!( + builder.skip_auto_import().type_signatures().build().snapshot(), + @"zqzqzq :: int | str", + ); + } + + #[test] + fn multiple_declarations_function_scope2() { + let builder = completion_test_builder( + "\ +def foo(): + zqzqzq: int = 1 + zqzq + zqzqzq: str = 'foo' + return +", + ); + // The type for `zqzqzq` *should* be `int`, but we consider all + // reachable declarations and bindings, which means we get a + // union of `int` and `str` here even though the `str` binding + // doesn't exist at the cursor position. + assert_snapshot!( + builder.skip_auto_import().type_signatures().build().snapshot(), + @"zqzqzq :: int | str", + ); + } + + #[test] + fn multiple_declarations_function_parameter() { + let builder = completion_test_builder( + "\ +from pathlib import Path +def f(zqzqzq: str): + zqzqzq: Path = Path(zqzqzq) + zqzq + return +", + ); + // The type for `zqzqzq` *should* be `Path`, but we consider all + // reachable declarations and bindings, which means we get a + // union of `str` and `Path` here even though the `str` binding + // isn't live at the cursor position. + assert_snapshot!( + builder.skip_auto_import().type_signatures().build().snapshot(), + @"zqzqzq :: str | Path", + ); + } + /// A way to create a simple single-file (named `main.py`) completion test /// builder. /// diff --git a/crates/ty_python_semantic/src/semantic_index/use_def.rs b/crates/ty_python_semantic/src/semantic_index/use_def.rs index 5683fc736b..dbd26595fd 100644 --- a/crates/ty_python_semantic/src/semantic_index/use_def.rs +++ b/crates/ty_python_semantic/src/semantic_index/use_def.rs @@ -590,6 +590,30 @@ impl<'db> UseDefMap<'db> { .map(|symbol_id| (symbol_id, self.end_of_scope_symbol_bindings(symbol_id))) } + pub(crate) fn all_reachable_symbols<'map>( + &'map self, + ) -> impl Iterator< + Item = ( + ScopedSymbolId, + DeclarationsIterator<'map, 'db>, + BindingWithConstraintsIterator<'map, 'db>, + ), + > + 'map { + self.reachable_definitions_by_symbol.iter_enumerated().map( + |(symbol_id, reachable_definitions)| { + let declarations = self.declarations_iterator( + &reachable_definitions.declarations, + BoundnessAnalysis::AssumeBound, + ); + let bindings = self.bindings_iterator( + &reachable_definitions.bindings, + BoundnessAnalysis::AssumeBound, + ); + (symbol_id, declarations, bindings) + }, + ) + } + /// This function is intended to be called only once inside `TypeInferenceBuilder::infer_function_body`. pub(crate) fn can_implicitly_return_none(&self, db: &dyn crate::Db) -> bool { !self diff --git a/crates/ty_python_semantic/src/semantic_model.rs b/crates/ty_python_semantic/src/semantic_model.rs index af887277d2..54ca0ba74f 100644 --- a/crates/ty_python_semantic/src/semantic_model.rs +++ b/crates/ty_python_semantic/src/semantic_model.rs @@ -11,7 +11,7 @@ use crate::module_resolver::{KnownModule, Module, list_modules, resolve_module}; use crate::semantic_index::definition::Definition; use crate::semantic_index::scope::FileScopeId; use crate::semantic_index::semantic_index; -use crate::types::list_members::{Member, all_end_of_scope_members, all_members}; +use crate::types::list_members::{Member, all_members, all_reachable_members}; use crate::types::{Type, binding_type, infer_scope_types}; use crate::{Db, resolve_real_shadowable_module}; @@ -76,7 +76,7 @@ impl<'db> SemanticModel<'db> { for (file_scope, _) in index.ancestor_scopes(file_scope) { for memberdef in - all_end_of_scope_members(self.db, file_scope.to_scope_id(self.db, self.file)) + all_reachable_members(self.db, file_scope.to_scope_id(self.db, self.file)) { members.insert( memberdef.member.name, @@ -221,7 +221,7 @@ impl<'db> SemanticModel<'db> { let mut completions = vec![]; for (file_scope, _) in index.ancestor_scopes(file_scope) { completions.extend( - all_end_of_scope_members(self.db, file_scope.to_scope_id(self.db, self.file)).map( + all_reachable_members(self.db, file_scope.to_scope_id(self.db, self.file)).map( |memberdef| Completion { name: memberdef.member.name, ty: Some(memberdef.member.ty), diff --git a/crates/ty_python_semantic/src/types/list_members.rs b/crates/ty_python_semantic/src/types/list_members.rs index 81ddac3d5b..a93438a596 100644 --- a/crates/ty_python_semantic/src/types/list_members.rs +++ b/crates/ty_python_semantic/src/types/list_members.rs @@ -25,7 +25,8 @@ use crate::{ }, }; -/// Iterate over all declarations and bindings in the given scope. +/// Iterate over all declarations and bindings that exist at the end +/// of the given scope. pub(crate) fn all_end_of_scope_members<'db>( db: &'db dyn Db, scope_id: ScopeId<'db>, @@ -75,6 +76,60 @@ pub(crate) fn all_end_of_scope_members<'db>( )) } +/// Iterate over all declarations and bindings that are reachable anywhere +/// in the given scope. +pub(crate) fn all_reachable_members<'db>( + db: &'db dyn Db, + scope_id: ScopeId<'db>, +) -> impl Iterator> + 'db { + let use_def_map = use_def_map(db, scope_id); + let table = place_table(db, scope_id); + + use_def_map + .all_reachable_symbols() + .flat_map(move |(symbol_id, declarations, bindings)| { + let symbol = table.symbol(symbol_id); + + let declaration_place_result = place_from_declarations(db, declarations); + let declaration = + declaration_place_result + .first_declaration + .and_then(|first_reachable_definition| { + let ty = declaration_place_result + .ignore_conflicting_declarations() + .place + .ignore_possibly_undefined()?; + let member = Member { + name: symbol.name().clone(), + ty, + }; + Some(MemberWithDefinition { + member, + first_reachable_definition, + }) + }); + + let place_with_definition = place_from_bindings(db, bindings); + let binding = + place_with_definition + .first_definition + .and_then(|first_reachable_definition| { + let ty = place_with_definition.place.ignore_possibly_undefined()?; + let member = Member { + name: symbol.name().clone(), + ty, + }; + Some(MemberWithDefinition { + member, + first_reachable_definition, + }) + }); + + [declaration, binding] + }) + .flatten() +} + // `__init__`, `__repr__`, `__eq__`, `__ne__` and `__hash__` are always included via `object`, // so we don't need to list them here. const SYNTHETIC_DATACLASS_ATTRIBUTES: &[&str] = &[