[ty] Fix missing registration of identifiers during semantic index construction

Fixes #572

_typos.toml

@@ -4,6 +4,10 @@ extend-exclude = [
     "crates/ty_vendored/vendor/**/*",
     "**/resources/**/*",
     "**/snapshots/**/*",
+    # Completion tests tend to have a lot of incomplete
+    # words naturally. It's annoying to have to make all
+    # of them actually words. So just ignore typos here.
+    "crates/ty_ide/src/completion.rs",
 ]
 
 [default.extend-words]

crates/ty_ide/src/completion.rs

@@ -861,6 +861,181 @@ print(f\"{some<CURSOR>
         ");
     }
 
+    // Ref: https://github.com/astral-sh/ty/issues/572
+    #[test]
+    fn scope_id_missing_function_identifier1() {
+        let test = cursor_test(
+            "\
+def m<CURSOR>
+",
+        );
+
+        assert_snapshot!(test.completions(), @r"
+        m
+        ");
+    }
+
+    // Ref: https://github.com/astral-sh/ty/issues/572
+    #[test]
+    fn scope_id_missing_function_identifier2() {
+        let test = cursor_test(
+            "\
+def m<CURSOR>(): pass
+",
+        );
+
+        assert_snapshot!(test.completions(), @r"
+        m
+        ");
+    }
+
+    // Ref: https://github.com/astral-sh/ty/issues/572
+    #[test]
+    fn fscope_id_missing_function_identifier3() {
+        let test = cursor_test(
+            "\
+def m(): pass
+<CURSOR>
+",
+        );
+
+        assert_snapshot!(test.completions(), @r"
+        m
+        ");
+    }
+
+    // Ref: https://github.com/astral-sh/ty/issues/572
+    #[test]
+    fn scope_id_missing_class_identifier1() {
+        let test = cursor_test(
+            "\
+class M<CURSOR>
+",
+        );
+
+        assert_snapshot!(test.completions(), @r"
+        M
+        ");
+    }
+
+    // Ref: https://github.com/astral-sh/ty/issues/572
+    #[test]
+    fn scope_id_missing_type_alias1() {
+        let test = cursor_test(
+            "\
+Fo<CURSOR> = float
+",
+        );
+
+        assert_snapshot!(test.completions(), @r"
+        Fo
+        float
+        ");
+    }
+
+    // Ref: https://github.com/astral-sh/ty/issues/572
+    #[test]
+    fn scope_id_missing_import1() {
+        let test = cursor_test(
+            "\
+import fo<CURSOR>
+",
+        );
+
+        assert_snapshot!(test.completions(), @r"
+        fo
+        ");
+    }
+
+    // Ref: https://github.com/astral-sh/ty/issues/572
+    #[test]
+    fn scope_id_missing_import2() {
+        let test = cursor_test(
+            "\
+import foo as ba<CURSOR>
+",
+        );
+
+        assert_snapshot!(test.completions(), @r"
+        ba
+        ");
+    }
+
+    // Ref: https://github.com/astral-sh/ty/issues/572
+    #[test]
+    fn scope_id_missing_from_import1() {
+        let test = cursor_test(
+            "\
+from fo<CURSOR> import wat
+",
+        );
+
+        assert_snapshot!(test.completions(), @r"
+        wat
+        ");
+    }
+
+    // Ref: https://github.com/astral-sh/ty/issues/572
+    #[test]
+    fn scope_id_missing_from_import2() {
+        let test = cursor_test(
+            "\
+from foo import wa<CURSOR>
+",
+        );
+
+        assert_snapshot!(test.completions(), @r"
+        wa
+        ");
+    }
+
+    // Ref: https://github.com/astral-sh/ty/issues/572
+    #[test]
+    fn scope_id_missing_from_import3() {
+        let test = cursor_test(
+            "\
+from foo import wat as ba<CURSOR>
+",
+        );
+
+        assert_snapshot!(test.completions(), @r"
+        ba
+        ");
+    }
+
+    // Ref: https://github.com/astral-sh/ty/issues/572
+    #[test]
+    fn scope_id_missing_try_except1() {
+        let test = cursor_test(
+            "\
+try:
+    pass
+except Type<CURSOR>:
+    pass
+",
+        );
+
+        assert_snapshot!(test.completions(), @r"
+        Type
+        ");
+    }
+
+    // Ref: https://github.com/astral-sh/ty/issues/572
+    #[test]
+    fn scope_id_missing_global1() {
+        let test = cursor_test(
+            "\
+def _():
+    global fo<CURSOR>
+",
+        );
+
+        assert_snapshot!(test.completions(), @r"
+        _
+        fo
+        ");
+    }
+
     impl CursorTest {
         fn completions(&self) -> String {
             let completions = completion(&self.db, self.file, self.cursor_offset);

crates/ty_python_semantic/src/semantic_index/builder.rs

@@ -1040,6 +1040,11 @@ impl<'db> SemanticIndexBuilder<'db> {
         }
     }
 
+    fn record_scope_for_identifier(&mut self, name: &ast::Identifier) {
+        self.scopes_by_expression
+            .insert(name.into(), self.current_scope());
+    }
+
     pub(super) fn build(mut self) -> SemanticIndex<'db> {
         let module = self.module;
         self.visit_body(module.suite());
@@ -1139,6 +1144,7 @@ where
                     is_async: _,
                     range: _,
                 } = function_def;
+                self.record_scope_for_identifier(name);
                 for decorator in decorator_list {
                     self.visit_decorator(decorator);
                 }
@@ -1221,6 +1227,7 @@ where
                 self.add_definition(symbol, function_def);
             }
             ast::Stmt::ClassDef(class) => {
+                self.record_scope_for_identifier(&class.name);
                 for decorator in &class.decorator_list {
                     self.visit_decorator(decorator);
                 }
@@ -1270,6 +1277,10 @@ where
                     .record_node_reachability(NodeKey::from_node(node));
 
                 for (alias_index, alias) in node.names.iter().enumerate() {
+                    self.record_scope_for_identifier(&alias.name);
+                    if let Some(ref asname) = alias.asname {
+                        self.record_scope_for_identifier(asname);
+                    }
                     // Mark the imported module, and all of its parents, as being imported in this
                     // file.
                     if let Some(module_name) = ModuleName::new(&alias.name) {
@@ -1294,6 +1305,9 @@ where
                 }
             }
             ast::Stmt::ImportFrom(node) => {
+                if let Some(ref module) = node.module {
+                    self.record_scope_for_identifier(module);
+                }
                 self.current_use_def_map_mut()
                     .record_node_reachability(NodeKey::from_node(node));
 
@@ -1378,6 +1392,10 @@ where
                         continue;
                     }
 
+                    self.record_scope_for_identifier(&alias.name);
+                    if let Some(ref asname) = alias.asname {
+                        self.record_scope_for_identifier(asname);
+                    }
                     let (symbol_name, is_reexported) = if let Some(asname) = &alias.asname {
                         (&asname.id, asname.id == alias.name.id)
                     } else {
@@ -1919,6 +1937,7 @@ where
             }
             ast::Stmt::Global(ast::StmtGlobal { range: _, names }) => {
                 for name in names {
+                    self.record_scope_for_identifier(name);
                     let symbol_id = self.add_symbol(name.id.clone());
                     let symbol_table = self.current_symbol_table();
                     let symbol = symbol_table.symbol(symbol_id);