fix(definition): import submodule as alias

pyrefly/lib/export/exports.rs

@@ -49,9 +49,9 @@ pub struct Export {
 /// Where is this export defined?
 #[derive(Debug, Clone)]
 pub enum ExportLocation {
-    // This export is defined in this module.
+    /// This export is defined in this module.
     ThisModule(Export),
-    // Export from another module ModuleName. If it's aliased, the old name (before the alias) is provided.
+    /// Export from another module ModuleName. If it's aliased, the old name (before the alias) is provided.
     OtherModule(ModuleName, Option<Name>),
 }
 

pyrefly/lib/state/lsp.rs

@@ -507,6 +507,20 @@ impl<'a> Transaction<'a> {
         ans.get_chosen_overload_trace(range)
     }
 
+    fn import_handle_with_preference(
+        &self,
+        handle: &Handle,
+        module: ModuleName,
+        preference: FindPreference,
+    ) -> Option<Handle> {
+        match preference.prefer_pyi {
+            true => self.import_handle(handle, module, None).finding(),
+            false => self
+                .import_handle_prefer_executable(handle, module, None)
+                .finding(),
+        }
+    }
+
     fn type_from_expression_at(&self, handle: &Handle, position: TextSize) -> Option<Type> {
         let module = self.get_ast(handle)?;
         let covering_nodes = Ast::locate_node(&module, position);
@@ -940,12 +954,7 @@ impl<'a> Transaction<'a> {
         let mut gas = RESOLVE_EXPORT_INITIAL_GAS;
         let mut name = name;
         while !gas.stop() {
-            let handle = match preference.prefer_pyi {
-                true => self.import_handle(handle, m, None).finding()?,
-                false => self
-                    .import_handle_prefer_executable(handle, m, None)
-                    .finding()?,
-            };
+            let handle = self.import_handle_with_preference(handle, m, preference)?;
             match self.get_exports(&handle).get(&name) {
                 Some(ExportLocation::ThisModule(export)) => {
                     return Some((handle.clone(), export.clone()));
@@ -954,6 +963,22 @@ impl<'a> Transaction<'a> {
                     if let Some(aliased_name) = aliased_name {
                         name = aliased_name.clone();
                     }
+                    if *module == m && handle.path().is_init() {
+                        let submodule = m.append(&name);
+                        let sub_handle =
+                            self.import_handle_with_preference(&handle, submodule, preference)?;
+                        let docstring_range = self.get_module_docstring_range(&sub_handle);
+                        return Some((
+                            sub_handle,
+                            Export {
+                                location: TextRange::default(),
+                                symbol_kind: Some(SymbolKind::Module),
+                                docstring_range,
+                                deprecation: None,
+                                special_export: None,
+                            },
+                        ));
+                    }
                     m = *module;
                 }
                 None => return None,
@@ -1023,12 +1048,7 @@ impl<'a> Transaction<'a> {
                         },
                     ));
                 }
-                let handle = match preference.prefer_pyi {
-                    true => self.import_handle(handle, name, None).finding()?,
-                    false => self
-                        .import_handle_prefer_executable(handle, name, None)
-                        .finding()?,
-                };
+                let handle = self.import_handle_with_preference(handle, name, preference)?;
                 let docstring_range = self.get_module_docstring_range(&handle);
                 Some((
                     handle,
@@ -1383,12 +1403,7 @@ impl<'a> Transaction<'a> {
         preference: FindPreference,
     ) -> Option<FindDefinitionItemWithDocstring> {
         // TODO: Handle relative import (via ModuleName::new_maybe_relative)
-        let handle = match preference.prefer_pyi {
-            true => self.import_handle(handle, module_name, None).finding()?,
-            false => self
-                .import_handle_prefer_executable(handle, module_name, None)
-                .finding()?,
-        };
+        let handle = self.import_handle_with_preference(handle, module_name, preference)?;
         // if the module is not yet loaded, force loading by asking for exports
         // necessary for imports that are not in tdeps (e.g. .py when there is also a .pyi)
         // todo(kylei): better solution

pyrefly/lib/test/lsp/definition.rs

@@ -1758,14 +1758,14 @@ Definition Result:
 }
 
 #[test]
-fn goto_def_on_import_different_name_alias_first_token_test() {
+fn goto_def_on_import_different_name_alias_test() {
     let lib = r#"
 def bar():
     pass
 "#;
     let code = r#"
 from lib import bar as baz
-#                ^
+#                ^      ^
 "#;
     let report =
         get_batched_lsp_operations_report(&[("main", code), ("lib", lib)], get_test_report);
@@ -1778,29 +1778,6 @@ Definition Result:
 2 | def bar():
         ^^^
 
-
-# lib.py
-"#
-        .trim(),
-        report.trim(),
-    );
-}
-
-#[test]
-fn goto_def_on_import_different_name_alias_second_token_test() {
-    let lib = r#"
-def bar():
-    pass
-"#;
-    let code = r#"
-from lib import bar as baz
-#                       ^
-"#;
-    let report =
-        get_batched_lsp_operations_report(&[("main", code), ("lib", lib)], get_test_report);
-    assert_eq!(
-        r#"
-# main.py
 2 | from lib import bar as baz
                             ^
 Definition Result:

pyrefly/lib/test/lsp/lsp_interaction/definition.rs

@@ -535,3 +535,20 @@ fn goto_def_on_none_goes_to_builtins_stub() {
         })
         .unwrap();
 }
+
+#[test]
+fn test_goto_def_imported_submodule_with_alias() {
+    let root = get_test_files_root();
+    let root_path = root.path().join("nested_package_imports");
+    test_go_to_def(
+        root_path,
+        None,
+        "main.py",
+        vec![
+            // `from pkg import sub as sub` -> first `sub`
+            (5, 16, "pkg/sub.py", 0, 0, 0, 0),
+            // `from pkg import sub as sub` -> second `sub`
+            (5, 23, "pkg/sub.py", 0, 0, 0, 0),
+        ],
+    );
+}

pyrefly/lib/test/lsp/lsp_interaction/semantic_tokens.rs

@@ -20,7 +20,7 @@ use crate::test::lsp::lsp_interaction::util::get_test_files_root;
 #[test]
 fn semantic_tokens_import_submodule_alias() {
     let root = get_test_files_root();
-    let root_path = root.path().join("semantic_tokens_imports");
+    let root_path = root.path().join("nested_package_imports");
     let mut interaction = LspInteraction::new();
     interaction.set_root(root_path.clone());
     interaction