implement basic completions for module attributes

Summary:
this diff implements basic completions for the `y` in `from x import y`.

we do this by reusing the `import_at` api implemented in the last diff. unfortunately, there are a few gaps:
- `from x import <>` does not work
- `from x imp<>` unfortunately works (because ruff's parser makes that an ImportFrom node with module = `x` and names = ['imp'])

I will fix these in follow-up diffs, but I believe the current implementation is better than nothing for now

Reviewed By: SamChou19815

Differential Revision: D75477356

fbshipit-source-id: 5fabe76f474a842f2d5e8319b50c7cc567de455d

Author: kinto0

pyrefly/lib/state/lsp.rs

@@ -551,7 +551,28 @@ impl<'a> Transaction<'a> {
         handle: &Handle,
         position: TextSize,
     ) -> Option<Vec<CompletionItem>> {
-        if self.identifier_at(handle, position).is_some() {
+        if let Some(import) = self.import_at(handle, position) {
+            return match import {
+                ImportIdentifier::Name(module_name) => {
+                    let handle = self.import_handle(handle, module_name, None).ok()?;
+                    let exports = self.get_exports(&handle);
+                    let completions = exports
+                        .keys()
+                        .map(|name| CompletionItem {
+                            label: name.to_string(),
+                            // todo(kylei): completion kind for exports
+                            kind: Some(CompletionItemKind::VARIABLE),
+                            ..Default::default()
+                        })
+                        .collect();
+                    Some(completions)
+                }
+                ImportIdentifier::Module(_module_name) => {
+                    // TODO(kylei): completion for module names
+                    None
+                }
+            };
+        } else if self.identifier_at(handle, position).is_some() {
             let bindings = self.get_bindings(handle)?;
             let module_info = self.get_module_info(handle)?;
             let names = bindings

pyrefly/lib/test/lsp/completion.rs

@@ -179,3 +179,94 @@ Completion Results:
         report.trim(),
     );
 }
+
+// TODO(kylei): ruff's ast gives us names = ["imp"] for `from foo imp`
+#[test]
+fn from_import_imp_test() {
+    let foo_code = r#"
+imperial_guard = "cool"
+"#;
+    let main_code = r#"
+from foo imp
+#          ^
+"#;
+    let report = get_batched_lsp_operations_report_allow_error(
+        &[("main", main_code), ("foo", foo_code)],
+        get_test_report,
+    );
+    assert_eq!(
+        r#"
+# main.py
+2 | from foo imp
+               ^
+Completion Results:
+- (Variable) imperial_guard
+
+
+# foo.py
+"#
+        .trim(),
+        report.trim(),
+    );
+}
+// TODO(kylei): ruff's ast gives us names = [] for `from foo import <>`
+#[test]
+fn from_import_empty_test() {
+    let foo_code = r#"
+imperial_guard = "cool"
+"#;
+    let main_code = r#"
+from foo import 
+#              ^
+"#;
+    let report = get_batched_lsp_operations_report_allow_error(
+        &[("main", main_code), ("foo", foo_code)],
+        get_test_report,
+    );
+    assert_eq!(
+        r#"
+# main.py
+2 | from foo import 
+                   ^
+Completion Results:
+
+
+# foo.py
+"#
+        .trim(),
+        report.trim(),
+    );
+}
+
+#[test]
+fn from_import_basic() {
+    let foo_code = r#"
+imperial_guard = "cool"
+"#;
+    let main_code = r#"
+from foo import imperial
+#          ^           ^
+"#;
+    let report = get_batched_lsp_operations_report_allow_error(
+        &[("main", main_code), ("foo", foo_code)],
+        get_test_report,
+    );
+    assert_eq!(
+        r#"
+# main.py
+2 | from foo import imperial
+               ^
+Completion Results:
+
+2 | from foo import imperial
+                           ^
+Completion Results:
+- (Variable) imperial_guard
+
+
+# foo.py
+"#
+        .trim(),
+        report.trim(),
+    );
+}