[ty] Adjust scope completions to use all reachable symbols

Fixes astral-sh/ty#1294

crates/ty_ide/src/completion.rs

@@ -6433,6 +6433,155 @@ collabc<CURSOR>
         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_<CURSOR>
+    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_<CURSOR>
+            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<CURSOR>
+",
+        );
+        // 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<CURSOR>
+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<CURSOR>
+    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<CURSOR>
+    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<CURSOR>
+    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.
     ///

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

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),

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<Item = MemberWithDefinition<'db>> + '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] = &[