diff --git a/crates/red_knot_python_semantic/src/module_resolver/path.rs b/crates/red_knot_python_semantic/src/module_resolver/path.rs index a49d1dc20d70c..0b8a2a26d9e4e 100644 --- a/crates/red_knot_python_semantic/src/module_resolver/path.rs +++ b/crates/red_knot_python_semantic/src/module_resolver/path.rs @@ -59,6 +59,10 @@ impl ModulePath { self.relative_path.push(component); } + pub(crate) fn pop(&mut self) -> bool { + self.relative_path.pop() + } + #[must_use] pub(super) fn is_directory(&self, resolver: &ResolverContext) -> bool { let ModulePath { diff --git a/crates/red_knot_python_semantic/src/module_resolver/resolver.rs b/crates/red_knot_python_semantic/src/module_resolver/resolver.rs index 5648dcd24fb80..daf25fd0704ec 100644 --- a/crates/red_knot_python_semantic/src/module_resolver/resolver.rs +++ b/crates/red_knot_python_semantic/src/module_resolver/resolver.rs @@ -569,24 +569,16 @@ fn resolve_name(db: &dyn Db, name: &ModuleName) -> Option<(SearchPath, File, Mod package_path.push(module_name); - // Must be a `__init__.pyi` or `__init__.py` or it isn't a package. - let kind = if package_path.is_directory(&resolver_state) { - package_path.push("__init__"); - ModuleKind::Package - } else { - ModuleKind::Module - }; - - // TODO Implement full https://peps.python.org/pep-0561/#type-checker-module-resolution-order resolution - if let Some(stub) = package_path.with_pyi_extension().to_file(&resolver_state) { - return Some((search_path.clone(), stub, kind)); + // Check for a regular package first (highest priority) + package_path.push("__init__"); + if let Some(regular_package) = resolve_file_module(&package_path, &resolver_state) { + return Some((search_path.clone(), regular_package, ModuleKind::Package)); } - if let Some(module) = package_path - .with_py_extension() - .and_then(|path| path.to_file(&resolver_state)) - { - return Some((search_path.clone(), module, kind)); + // Check for a file module next + package_path.pop(); + if let Some(file_module) = resolve_file_module(&package_path, &resolver_state) { + return Some((search_path.clone(), file_module, ModuleKind::Module)); } // For regular packages, don't search the next search path. All files of that @@ -607,6 +599,23 @@ fn resolve_name(db: &dyn Db, name: &ModuleName) -> Option<(SearchPath, File, Mod None } +/// If `module` exists on disk with either a `.pyi` or `.py` extension, +/// return the [`File`] corresponding to that path. +/// +/// `.pyi` files take priority, as they always have priority when +/// resolving modules. +fn resolve_file_module(module: &ModulePath, resolver_state: &ResolverContext) -> Option { + // Stubs have precedence over source files + module + .with_pyi_extension() + .to_file(resolver_state) + .or_else(|| { + module + .with_py_extension() + .and_then(|path| path.to_file(resolver_state)) + }) +} + fn resolve_package<'a, 'db, I>( module_search_path: &SearchPath, components: I, @@ -633,7 +642,10 @@ where if is_regular_package { in_namespace_package = false; - } else if package_path.is_directory(resolver_state) { + } else if package_path.is_directory(resolver_state) + // Pure modules hide namespace packages with the same name + && resolve_file_module(&package_path, resolver_state).is_none() + { // A directory without an `__init__.py` is a namespace package, continue with the next folder. in_namespace_package = true; } else if in_namespace_package { @@ -1091,6 +1103,25 @@ mod tests { ); } + #[test] + fn single_file_takes_priority_over_namespace_package() { + //const SRC: &[FileSpec] = &[("foo.py", "x = 1")]; + const SRC: &[FileSpec] = &[("foo.py", "x = 1"), ("foo/bar.py", "x = 2")]; + + let TestCase { db, src, .. } = TestCaseBuilder::new().with_src_files(SRC).build(); + + let foo_module_name = ModuleName::new_static("foo").unwrap(); + let foo_bar_module_name = ModuleName::new_static("foo.bar").unwrap(); + + // `foo.py` takes priority over the `foo` namespace package + let foo_module = resolve_module(&db, foo_module_name.clone()).unwrap(); + assert_eq!(foo_module.file().path(&db), &src.join("foo.py")); + + // `foo.bar` isn't recognised as a module + let foo_bar_module = resolve_module(&db, foo_bar_module_name.clone()); + assert_eq!(foo_bar_module, None); + } + #[test] fn typing_stub_over_module() { const SRC: &[FileSpec] = &[("foo.py", "print('Hello, world!')"), ("foo.pyi", "x: int")];