From 2022bebfb6cb6a33eb7c5c6846c154ac5b14c68a Mon Sep 17 00:00:00 2001 From: Ting-Hong Shieh <32212900+ting-hong-shieh@users.noreply.github.com> Date: Thu, 13 Aug 2026 16:24:32 +0800 Subject: [PATCH] Preserve typed overlays for editable packages Keep generated stubs from typed namespace roots available when a later editable source directory provides the regular package initializer. --- crates/pyrefly_build/src/module_resolver.rs | 82 ++++++++++++++++++++- pyrefly/lib/module/finder.rs | 54 ++++++++++++++ pyrefly/lib/test/imports.rs | 39 ++++++++++ 3 files changed, 172 insertions(+), 3 deletions(-) diff --git a/crates/pyrefly_build/src/module_resolver.rs b/crates/pyrefly_build/src/module_resolver.rs index b8a599dc7c..170a3aa32b 100644 --- a/crates/pyrefly_build/src/module_resolver.rs +++ b/crates/pyrefly_build/src/module_resolver.rs @@ -233,6 +233,10 @@ pub enum FindResult { /// Found a regular package. The first field is its `__init__` file, and the /// second is the package directory used as the sole root for submodule lookup. RegularPackage(PathBuf, PathBuf), + /// Found a regular package after one or more typed namespace directories. + /// The namespace directories contain `py.typed` and provide overlays for + /// generated extension stubs in editable installs. + RegularPackageWithTypedOverlay(PathBuf, Vec1), /// Found a legacy namespace package: a package whose `__init__` calls /// `pkgutil.extend_path`. The first field is the winning `__init__`; the /// roots accumulate same-named directories across the search path. @@ -265,8 +269,18 @@ impl FindResult { // RegularPackage and LegacyNamespacePackage share the top tier: both // resolve to a concrete `__init__`. Tying them lets the prefer-`a` // rule preserve sys.path order when fallback roots are folded in. - (FindResult::RegularPackage(..), _) | (FindResult::LegacyNamespacePackage(..), _) => a, - (_, FindResult::RegularPackage(..)) | (_, FindResult::LegacyNamespacePackage(..)) => b, + ( + FindResult::RegularPackage(..) + | FindResult::RegularPackageWithTypedOverlay(..) + | FindResult::LegacyNamespacePackage(..), + _, + ) => a, + ( + _, + FindResult::RegularPackage(..) + | FindResult::RegularPackageWithTypedOverlay(..) + | FindResult::LegacyNamespacePackage(..), + ) => b, (FindResult::SingleFilePyiModule(_), _) => a, (_, FindResult::SingleFilePyiModule(_)) => b, (FindResult::SingleFilePyModule(_), _) => a, @@ -286,6 +300,9 @@ impl FindResult { FindResult::RegularPackage(_, package_dir) => { dir_cache.is_partial_py_typed(&package_dir.join("py.typed"), observer) } + FindResult::RegularPackageWithTypedOverlay(_, package_dirs) => package_dirs + .iter() + .any(|dir| dir_cache.is_partial_py_typed(&dir.join("py.typed"), observer)), FindResult::LegacyNamespacePackage(init_path, _) => { init_path.parent().is_some_and(|package_dir| { dir_cache.is_partial_py_typed(&package_dir.join("py.typed"), observer) @@ -307,6 +324,7 @@ impl FindResult { FindResult::SingleFilePyiModule(path) | FindResult::SingleFilePyModule(path) | FindResult::RegularPackage(path, _) + | FindResult::RegularPackageWithTypedOverlay(path, _) | FindResult::LegacyNamespacePackage(path, _) => Some(ModulePath::filesystem(path)), FindResult::ImplicitNamespacePackage(roots) => { Some(ModulePath::namespace(roots.first().clone())) @@ -330,6 +348,13 @@ pub fn package_has_py_typed( let depth = module.components().len().saturating_sub(1); let (dir, levels_up) = match result { FindResult::RegularPackage(_, dir) => (dir.as_path(), depth), + FindResult::RegularPackageWithTypedOverlay(_, dirs) => { + return dirs.iter().any(|dir| { + iter::successors(Some(dir.as_path()), |dir| dir.parent()) + .take(depth + 1) + .any(|dir| dir_cache.file_exists(&dir.join("py.typed"))) + }); + } FindResult::LegacyNamespacePackage(init_path, _) => { let Some(dir) = init_path.parent() else { return false; @@ -501,7 +526,27 @@ fn find_one_part<'a>( } } Some(FindResult::RegularPackage(init_path, init_dir)) => match &mut acc { - None | Some(NamespaceAccumulator::Implicit(_)) => { + Some(NamespaceAccumulator::Implicit(namespace_roots)) => { + let typed_overlays = namespace_roots + .iter() + .filter(|root| { + timed_stat(observer, || dir_cache.file_exists(&root.join("py.typed"))) + }) + .cloned() + .collect::>(); + if let Ok(mut package_roots) = Vec1::try_from_vec(typed_overlays) { + package_roots.push(init_dir); + return Some(( + FindResult::RegularPackageWithTypedOverlay(init_path, package_roots), + roots.cloned().collect(), + )); + } + return Some(( + FindResult::RegularPackage(init_path, init_dir), + roots.cloned().collect(), + )); + } + None => { return Some(( FindResult::RegularPackage(init_path, init_dir), roots.cloned().collect(), @@ -551,6 +596,17 @@ fn continue_find_module( ) .map(|x| x.0); } + Some(FindResult::RegularPackageWithTypedOverlay(_, next_roots)) => { + current_result = find_one_part( + part, + next_roots.iter(), + style_filter, + phantom_paths, + dir_cache, + observer, + ) + .map(|x| x.0); + } Some(FindResult::LegacyNamespacePackage(_, next_roots)) | Some(FindResult::ImplicitNamespacePackage(next_roots)) => { current_result = find_one_part( @@ -601,6 +657,7 @@ where match current_result { FindResult::SingleFilePyiModule(_) | FindResult::RegularPackage(..) + | FindResult::RegularPackageWithTypedOverlay(..) | FindResult::LegacyNamespacePackage(..) => Some(current_result), _ => Some( fallback_search @@ -800,6 +857,9 @@ fn find_one_part_prefix<'a>( if !results.iter().any(|r| match r { (FindResult::RegularPackage(_, p), _) => *p == path, + (FindResult::RegularPackageWithTypedOverlay(_, ps), _) => { + ps.iter().any(|p| *p == path) + } (FindResult::LegacyNamespacePackage(_, ps), _) => ps.first() == &path, _ => false, }) { @@ -883,6 +943,22 @@ pub fn find_module_prefixes<'a>( .map(|x| x.0); } } + Some(FindResult::RegularPackageWithTypedOverlay(_, next_roots)) => { + if is_last { + results = find_one_part_prefix(part, next_roots.iter()); + break; + } else { + current_result = find_one_part( + part, + next_roots.iter(), + None, + &mut None, + &dir_cache, + None, + ) + .map(|x| x.0); + } + } Some(FindResult::LegacyNamespacePackage(_, next_roots)) | Some(FindResult::ImplicitNamespacePackage(next_roots)) => { if is_last { diff --git a/pyrefly/lib/module/finder.rs b/pyrefly/lib/module/finder.rs index 3b513a39f1..5d8a68c139 100644 --- a/pyrefly/lib/module/finder.rs +++ b/pyrefly/lib/module/finder.rs @@ -1353,6 +1353,60 @@ mod tests { ); } + #[test] + fn test_typed_namespace_overlays_later_regular_package() { + let tempdir = tempfile::tempdir().unwrap(); + let root = tempdir.path(); + TestPath::setup_test_directory( + root, + vec![ + TestPath::dir( + "site_packages", + vec![TestPath::dir( + "a", + vec![TestPath::file("py.typed"), TestPath::file("b.pyi")], + )], + ), + TestPath::dir( + "editable_source", + vec![TestPath::dir("a", vec![TestPath::file("__init__.py")])], + ), + ], + ); + let roots = [root.join("site_packages"), root.join("editable_source")]; + + assert_eq!( + find_module( + ModuleName::from_str("a"), + roots.iter(), + &mut vec![], + None, + SitePackagePolicy::default(), + &mut None, + &DirEntryCache::new(), + None, + ) + .unwrap(), + FindingOrError::new_finding(ModulePath::filesystem( + root.join("editable_source/a/__init__.py") + )) + ); + assert_eq!( + find_module( + ModuleName::from_str("a.b"), + roots.iter(), + &mut vec![], + None, + SitePackagePolicy::default(), + &mut None, + &DirEntryCache::new(), + None, + ) + .unwrap(), + FindingOrError::new_finding(ModulePath::filesystem(root.join("site_packages/a/b.pyi"))) + ); + } + #[test] fn test_legacy_namespace_package_then_regular_package() { // Once `find_one_part` enters LegacyNamespacePackage (LNP) mode diff --git a/pyrefly/lib/test/imports.rs b/pyrefly/lib/test/imports.rs index 0bd6d17297..79eaf20d50 100644 --- a/pyrefly/lib/test/imports.rs +++ b/pyrefly/lib/test/imports.rs @@ -1986,6 +1986,45 @@ fn test_pkgutil_namespace_absorbs_implicit_namespace() { .unwrap(); } +#[test] +fn test_editable_install_generated_stub_overlay() { + let tempdir = tempfile::tempdir().unwrap(); + let root = tempdir.path(); + std::fs::create_dir_all(root.join("site_packages/package")).unwrap(); + std::fs::create_dir_all(root.join("editable_source/package")).unwrap(); + std::fs::write(root.join("site_packages/package/py.typed"), "").unwrap(); + std::fs::write( + root.join("site_packages/package/_core.pyi"), + "def add(a: int, b: int) -> int: ...\n", + ) + .unwrap(); + std::fs::write( + root.join("editable_source/package/__init__.py"), + "from package._core import add\n", + ) + .unwrap(); + + let mut env = TestEnv::new().with_site_package_paths(vec![ + root.join("site_packages"), + root.join("editable_source"), + ]); + env.add_with_path( + "main", + "main.py", + r#" +from package import add + +add("hello", "world") # E: Argument `Literal['hello']` is not assignable to parameter `a` with type `int` # E: Argument `Literal['world']` is not assignable to parameter `b` with type `int` +"#, + ); + let (state, handle_fn) = env.to_state(); + state + .transaction() + .get_errors(&[handle_fn("main")]) + .check_against_expectations() + .unwrap(); +} + // ---------------------------------------------------------------------------- // Cross-module class rebind tests: importers should observe whichever class // the visible result chose. See `assign.rs` for the same-module regressions.