represent imports better

Summary:
by using `Ast::imports` to find imports, we considered the entire import stmt including keywords in what we consider a "module". But in reality, we should only consider the names*.

This diff makes a `ImportIdentifier` enum, walking the statements and grabbing the necessary items.

For LSP features like completion, I believe the x in `from x import y` is very similar to the the x in `import x`. For that reason I make both of these `x`s `ImportIdentifier::Module`. But the names are different and get their own enum variant (to be used in next diff).

Reviewed By: SamChou19815

Differential Revision: D75473401

fbshipit-source-id: ddcf1b5da352e7bb9ab59e046784f09587f1cd02

Author: kinto0

pyrefly/lib/state/lsp.rs

@@ -60,6 +60,23 @@ pub enum DefinitionMetadata {
     VariableOrAttribute(Name),
 }
 
+enum ImportIdentifier {
+    // The name of a module. ex: `x` in `import x` or `from x import name`
+    Module(ModuleName),
+    // A name from a module's exports. ex: `name` in `from x import name`
+    // Note: these are also definitions
+    Name(ModuleName),
+}
+
+impl ImportIdentifier {
+    fn module_name(&self) -> ModuleName {
+        match self {
+            ImportIdentifier::Module(module_name) => *module_name,
+            ImportIdentifier::Name(module_name) => *module_name,
+        }
+    }
+}
+
 impl<'a> Transaction<'a> {
     fn get_type(&self, handle: &Handle, key: &Key) -> Option<Type> {
         let idx = self.get_bindings(handle)?.key_to_idx(key);
@@ -88,15 +105,43 @@ impl<'a> Transaction<'a> {
         res
     }
 
-    fn import_at(&self, handle: &Handle, position: TextSize) -> Option<ModuleName> {
-        let module = self.get_ast(handle)?;
-        for (module, text_range) in Ast::imports(&module, handle.module(), handle.path().is_init())
-        {
-            if text_range.contains_inclusive(position) {
-                return Some(module);
+    fn import_at(&self, handle: &Handle, position: TextSize) -> Option<ImportIdentifier> {
+        fn visit_stmt(x: &Stmt, find: TextSize, res: &mut Option<ImportIdentifier>) {
+            match x {
+                Stmt::Import(stmt_import) => {
+                    let mut parts = Vec::new();
+                    for name in stmt_import.names.iter() {
+                        parts.push(name.name.clone());
+                        if name.range.contains_inclusive(find) {
+                            *res = Some(ImportIdentifier::Module(ModuleName::from_parts(
+                                parts.clone(),
+                            )));
+                        }
+                    }
+                }
+                Stmt::ImportFrom(stmt_import_from) => {
+                    if let Some(id) = &stmt_import_from.module {
+                        if id.range.contains_inclusive(find) {
+                            *res = Some(ImportIdentifier::Module(ModuleName::from_name(&id.id)));
+                        } else {
+                            for name in stmt_import_from.names.iter() {
+                                if name.range.contains_inclusive(find) {
+                                    *res =
+                                        Some(ImportIdentifier::Name(ModuleName::from_name(&id.id)));
+                                }
+                            }
+                        }
+                    }
+                }
+                _ => x.recurse(&mut |x| visit_stmt(x, find, res)),
             }
         }
-        None
+
+        let mut res = None;
+        self.get_ast(handle)?
+            .body
+            .visit(&mut |x| visit_stmt(x, position, &mut res));
+        res
     }
 
     fn definition_at(&self, handle: &Handle, position: TextSize) -> Option<Key> {
@@ -133,9 +178,10 @@ impl<'a> Transaction<'a> {
             }
         }
         if let Some(m) = self.import_at(handle, position) {
+            let module_name = m.module_name();
             return Some(Type::Module(Module::new(
-                m.components().first().unwrap().clone(),
-                OrderedSet::from_iter([(m)]),
+                module_name.components().first().unwrap().clone(),
+                OrderedSet::from_iter([module_name]),
             )));
         }
         let attribute = self.attribute_at(handle, position)?;
@@ -273,7 +319,8 @@ impl<'a> Transaction<'a> {
             ));
         }
         if let Some(m) = self.import_at(handle, position) {
-            let handle = self.import_handle(handle, m, None).ok()?;
+            let module_name = m.module_name();
+            let handle = self.import_handle(handle, module_name, None).ok()?;
             return Some((
                 DefinitionMetadata::Module,
                 TextRangeWithModuleInfo::new(self.get_module_info(&handle)?, TextRange::default()),

pyrefly/lib/test/lsp/definition.rs

@@ -107,7 +107,7 @@ def f(x: list[int], y: str, z: Literal[42]):
     let code_test: &str = r#"
 from typing import Literal
 from .import_provider import f
-#           ^                ^
+# ^         ^                ^
 
 foo: Literal[1] = 1
 #        ^
@@ -125,6 +125,10 @@ bar = f([1], "", 42)
     assert_eq!(
         r#"
 # main.py
+3 | from .import_provider import f
+      ^
+Definition Result: None
+
 3 | from .import_provider import f
                 ^
 Definition Result:

pyrefly/lib/test/lsp/hover_type.rs

@@ -24,7 +24,7 @@ fn get_test_report(state: &State, handle: &Handle, position: TextSize) -> String
 fn basic_test() {
     let code = r#"
 from typing import Literal
-#        ^
+#        ^     ^        ^
 def f(x: list[int], y: str, z: Literal[42]):
 #   ^               ^       ^
     return x
@@ -40,6 +40,14 @@ yyy = f([1, 2, 3], "test", 42)
              ^
 Hover Result: `Module[typing]`
 
+2 | from typing import Literal
+                   ^
+Hover Result: None
+
+2 | from typing import Literal
+                            ^
+Hover Result: `type[Literal]`
+
 4 | def f(x: list[int], y: str, z: Literal[42]):
         ^
 Hover Result: `(x: list[int], y: str, z: Literal[42]) -> list[int]`
@@ -83,6 +91,25 @@ Hover Result: `Module[typing]`
     );
 }
 
+#[test]
+fn import_alias_test() {
+    let code = r#"
+import typing as t
+#                ^
+"#;
+    let report = get_batched_lsp_operations_report(&[("main", code)], get_test_report);
+    assert_eq!(
+        r#"
+# main.py
+2 | import typing as t
+                     ^
+Hover Result: `Module[typing]`
+"#
+        .trim(),
+        report.trim(),
+    );
+}
+
 #[test]
 fn duplicate_import_test() {
     let code = r#"
@@ -91,13 +118,12 @@ import typing
 #        ^
 "#;
     let report = get_batched_lsp_operations_report(&[("main", code)], get_test_report);
-    // TODO(kylei): The result should be `Module[typing]`
     assert_eq!(
         r#"
 # main.py
 3 | import typing
              ^
-Hover Result: None
+Hover Result: `Module[typing]`
 "#
         .trim(),
         report.trim(),