diff --git a/pyrefly/lib/module/bundled.rs b/pyrefly/lib/module/bundled.rs index e27273daed..773739f699 100644 --- a/pyrefly/lib/module/bundled.rs +++ b/pyrefly/lib/module/bundled.rs @@ -34,6 +34,17 @@ pub fn set_readonly(path: &Path, value: bool) -> anyhow::Result<()> { Ok(()) } +/// Prefer a package initializer when a bundle contains both `foo.pyi` and +/// `foo/__init__.pyi`, matching Python's filesystem import precedence. +pub(crate) fn bundled_module_path_is_preferred(candidate: &Path, existing: &Path) -> bool { + candidate + .file_name() + .is_some_and(|name| name == "__init__.pyi") + && existing + .file_name() + .is_none_or(|name| name != "__init__.pyi") +} + /// Creates a base config file for bundled stubs with common settings. /// /// This helper function encapsulates the common configuration logic shared across diff --git a/pyrefly/lib/module/third_party.rs b/pyrefly/lib/module/third_party.rs index 490ea9d3ba..a3a29f5eb6 100644 --- a/pyrefly/lib/module/third_party.rs +++ b/pyrefly/lib/module/third_party.rs @@ -21,6 +21,7 @@ use pyrefly_util::arc_id::ArcId; use starlark_map::small_map::SmallMap; use crate::module::bundled::BundledStub; +use crate::module::bundled::bundled_module_path_is_preferred; use crate::module::bundled::create_bundled_stub_config; #[derive(Debug, Clone)] @@ -60,7 +61,13 @@ impl BundledStub for BundledThirdParty { for (relative_path, contents) in contents { let adjusted_path = strip_stubs_suffix_from_path(&relative_path); let module_name = ModuleName::from_relative_path(&adjusted_path)?; - res.find.insert(module_name, relative_path.clone()); + if res + .find + .get(&module_name) + .is_none_or(|existing| bundled_module_path_is_preferred(&relative_path, existing)) + { + res.find.insert(module_name, relative_path.clone()); + } res.load.insert(relative_path, Arc::new(contents)); } Ok(res) @@ -140,6 +147,16 @@ mod tests { } } + #[test] + fn test_bundled_third_party_prefers_package_over_module() { + let stubs = get_bundled_third_party().unwrap(); + let path = stubs + .find + .get(&ModuleName::from_str("pandas.io.parsers")) + .expect("pandas.io.parsers should be bundled"); + assert_eq!(path.file_name().unwrap(), "__init__.pyi"); + } + #[test] fn test_bundled_third_party_load_works() { let stub = get_bundled_third_party().unwrap(); diff --git a/pyrefly/lib/module/typeshed.rs b/pyrefly/lib/module/typeshed.rs index 4f5d884805..c44dabb772 100644 --- a/pyrefly/lib/module/typeshed.rs +++ b/pyrefly/lib/module/typeshed.rs @@ -24,6 +24,7 @@ use starlark_map::small_map::SmallMap; use crate::config::config::ConfigFile; use crate::module::bundled::BundledStub; +use crate::module::bundled::bundled_module_path_is_preferred; use crate::module::bundled::create_bundled_stub_config; #[derive(Debug, Clone)] @@ -41,7 +42,13 @@ impl BundledStub for BundledTypeshedStdlib { }; for (relative_path, contents) in contents { let module_name = ModuleName::from_relative_path(&relative_path)?; - res.find.insert(module_name, relative_path.clone()); + if res + .find + .get(&module_name) + .is_none_or(|existing| bundled_module_path_is_preferred(&relative_path, existing)) + { + res.find.insert(module_name, relative_path.clone()); + } res.load.insert(relative_path, Arc::new(contents)); } Ok(res) diff --git a/pyrefly/lib/module/typeshed_third_party.rs b/pyrefly/lib/module/typeshed_third_party.rs index 965efa6107..af1c5474bd 100644 --- a/pyrefly/lib/module/typeshed_third_party.rs +++ b/pyrefly/lib/module/typeshed_third_party.rs @@ -20,6 +20,7 @@ use pyrefly_util::arc_id::ArcId; use starlark_map::small_map::SmallMap; use crate::module::bundled::BundledStub; +use crate::module::bundled::bundled_module_path_is_preferred; use crate::module::bundled::create_bundled_stub_config; #[derive(Debug, Clone)] @@ -45,8 +46,12 @@ impl BundledStub for BundledTypeshedThirdParty { .get(&relative_path) .map(|s| ModuleName::from_str(s)) .unwrap_or_else(|| ModuleName::from_name(&module_name.first_component())); - res.find - .insert(module_name, (relative_path.clone(), package_name)); + if res.find.get(&module_name).is_none_or(|(existing, _)| { + bundled_module_path_is_preferred(&relative_path, existing) + }) { + res.find + .insert(module_name, (relative_path.clone(), package_name)); + } res.load.insert(relative_path, Arc::new(contents)); } Ok(res) diff --git a/pyrefly/lib/test/lsp/lsp_interaction/hover.rs b/pyrefly/lib/test/lsp/lsp_interaction/hover.rs index 2a9cbbf684..e6421c67aa 100644 --- a/pyrefly/lib/test/lsp/lsp_interaction/hover.rs +++ b/pyrefly/lib/test/lsp/lsp_interaction/hover.rs @@ -115,6 +115,43 @@ fn hover_shows_third_party_function_name() { interaction.shutdown().unwrap(); } +#[test] +fn hover_shows_bundled_pandas_function_type() { + let root = get_test_files_root(); + let mut interaction = LspInteraction::new(); + interaction.set_root(root.path().join("pandas_read_csv_hover")); + interaction + .initialize(InitializeSettings { + configuration: Some(None), + ..Default::default() + }) + .unwrap(); + + interaction.client.did_open("test.py"); + interaction + .client + .definition("test.py", 2, 5) + .expect_definition_response_from_root( + "site_packages/pandas/io/parsers/readers.py", + 0, + 4, + 0, + 12, + ) + .unwrap(); + interaction + .client + .hover("test.py", 2, 5) + .expect_hover_response_with_markup(|value| { + value.is_some_and(|text| { + text.contains("read_csv:") && !text.contains("read_csv: Unknown") + }) + }) + .unwrap(); + + interaction.shutdown().unwrap(); +} + #[test] fn test_hover_import() { let root = get_test_files_root(); diff --git a/pyrefly/lib/test/lsp/lsp_interaction/test_files/pandas_read_csv_hover/pyrefly.toml b/pyrefly/lib/test/lsp/lsp_interaction/test_files/pandas_read_csv_hover/pyrefly.toml new file mode 100644 index 0000000000..c8a839ff11 --- /dev/null +++ b/pyrefly/lib/test/lsp/lsp_interaction/test_files/pandas_read_csv_hover/pyrefly.toml @@ -0,0 +1,2 @@ +search-path = ["."] +site-package-path = ["site_packages"] diff --git a/pyrefly/lib/test/lsp/lsp_interaction/test_files/pandas_read_csv_hover/site_packages/pandas/__init__.py b/pyrefly/lib/test/lsp/lsp_interaction/test_files/pandas_read_csv_hover/site_packages/pandas/__init__.py new file mode 100644 index 0000000000..fde53c955a --- /dev/null +++ b/pyrefly/lib/test/lsp/lsp_interaction/test_files/pandas_read_csv_hover/site_packages/pandas/__init__.py @@ -0,0 +1 @@ +from pandas.io.api import read_csv as read_csv diff --git a/pyrefly/lib/test/lsp/lsp_interaction/test_files/pandas_read_csv_hover/site_packages/pandas/io/__init__.py b/pyrefly/lib/test/lsp/lsp_interaction/test_files/pandas_read_csv_hover/site_packages/pandas/io/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/pyrefly/lib/test/lsp/lsp_interaction/test_files/pandas_read_csv_hover/site_packages/pandas/io/api.py b/pyrefly/lib/test/lsp/lsp_interaction/test_files/pandas_read_csv_hover/site_packages/pandas/io/api.py new file mode 100644 index 0000000000..8dea673cd9 --- /dev/null +++ b/pyrefly/lib/test/lsp/lsp_interaction/test_files/pandas_read_csv_hover/site_packages/pandas/io/api.py @@ -0,0 +1 @@ +from pandas.io.parsers.readers import read_csv as read_csv diff --git a/pyrefly/lib/test/lsp/lsp_interaction/test_files/pandas_read_csv_hover/site_packages/pandas/io/parsers/__init__.py b/pyrefly/lib/test/lsp/lsp_interaction/test_files/pandas_read_csv_hover/site_packages/pandas/io/parsers/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/pyrefly/lib/test/lsp/lsp_interaction/test_files/pandas_read_csv_hover/site_packages/pandas/io/parsers/readers.py b/pyrefly/lib/test/lsp/lsp_interaction/test_files/pandas_read_csv_hover/site_packages/pandas/io/parsers/readers.py new file mode 100644 index 0000000000..affe70143a --- /dev/null +++ b/pyrefly/lib/test/lsp/lsp_interaction/test_files/pandas_read_csv_hover/site_packages/pandas/io/parsers/readers.py @@ -0,0 +1,2 @@ +def read_csv(path: str) -> object: + return object() diff --git a/pyrefly/lib/test/lsp/lsp_interaction/test_files/pandas_read_csv_hover/test.py b/pyrefly/lib/test/lsp/lsp_interaction/test_files/pandas_read_csv_hover/test.py new file mode 100644 index 0000000000..fb48cf7605 --- /dev/null +++ b/pyrefly/lib/test/lsp/lsp_interaction/test_files/pandas_read_csv_hover/test.py @@ -0,0 +1,3 @@ +import pandas as pd + +pd.read_csv("data.csv")