diff --git a/crates/red_knot_python_semantic/resources/mdtest/import/stub_packages.md b/crates/red_knot_python_semantic/resources/mdtest/import/stub_packages.md new file mode 100644 index 0000000000000..fef47bf47deb8 --- /dev/null +++ b/crates/red_knot_python_semantic/resources/mdtest/import/stub_packages.md @@ -0,0 +1,286 @@ +# Stub packages + +Stub packages are packages named `-stubs` that provide typing stubs for ``. See +[specification](https://typing.python.org/en/latest/spec/distributing.html#stub-only-packages). + +## Simple stub + +```toml +[environment] +extra-paths = ["/packages"] +``` + +`/packages/foo-stubs/__init__.pyi`: + +```pyi +class Foo: + name: str + age: int +``` + +`/packages/foo/__init__.py`: + +```py +class Foo: ... +``` + +`main.py`: + +```py +from foo import Foo + +reveal_type(Foo().name) # revealed: str +``` + +## Stubs only + +The regular package isn't required for type checking. + +```toml +[environment] +extra-paths = ["/packages"] +``` + +`/packages/foo-stubs/__init__.pyi`: + +```pyi +class Foo: + name: str + age: int +``` + +`main.py`: + +```py +from foo import Foo + +reveal_type(Foo().name) # revealed: str +``` + +## `-stubs` named module + +A module named `-stubs` isn't a stub package. + +```toml +[environment] +extra-paths = ["/packages"] +``` + +`/packages/foo-stubs.pyi`: + +```pyi +class Foo: + name: str + age: int +``` + +`main.py`: + +```py +from foo import Foo # error: [unresolved-import] + +reveal_type(Foo().name) # revealed: Unknown +``` + +## Namespace package in different search paths + +A namespace package with multiple stub packages spread over multiple search paths. + +```toml +[environment] +extra-paths = ["/stubs1", "/stubs2", "/packages"] +``` + +`/stubs1/shapes-stubs/polygons/pentagon.pyi`: + +```pyi +class Pentagon: + sides: int + area: float +``` + +`/stubs2/shapes-stubs/polygons/hexagon.pyi`: + +```pyi +class Hexagon: + sides: int + area: float +``` + +`/packages/shapes/polygons/pentagon.py`: + +```py +class Pentagon: ... +``` + +`/packages/shapes/polygons/hexagon.py`: + +```py +class Hexagon: ... +``` + +`main.py`: + +```py +from shapes.polygons.hexagon import Hexagon +from shapes.polygons.pentagon import Pentagon + +reveal_type(Pentagon().sides) # revealed: int +reveal_type(Hexagon().area) # revealed: int | float +``` + +## Inconsistent stub packages + +Stub packages where one is a namespace package and the other is a regular package. Module resolution +should stop after the first non-namespace stub package. This matches Pyright's behavior. + +```toml +[environment] +extra-paths = ["/stubs1", "/stubs2", "/packages"] +``` + +`/stubs1/shapes-stubs/__init__.pyi`: + +```pyi +``` + +`/stubs1/shapes-stubs/polygons/__init__.pyi`: + +```pyi +``` + +`/stubs1/shapes-stubs/polygons/pentagon.pyi`: + +```pyi +class Pentagon: + sides: int + area: float +``` + +`/stubs2/shapes-stubs/polygons/hexagon.pyi`: + +```pyi +class Hexagon: + sides: int + area: float +``` + +`/packages/shapes/polygons/pentagon.py`: + +```py +class Pentagon: ... +``` + +`/packages/shapes/polygons/hexagon.py`: + +```py +class Hexagon: ... +``` + +`main.py`: + +```py +from shapes.polygons.pentagon import Pentagon +from shapes.polygons.hexagon import Hexagon # error: [unresolved-import] + +reveal_type(Pentagon().sides) # revealed: int +reveal_type(Hexagon().area) # revealed: Unknown +``` + +## Namespace stubs for non-namespace package + +The runtime package is a regular package but the stubs are namespace packages. Pyright skips the +stub package if the "regular" package isn't a namespace package. I'm not aware that the behavior +here is specified, and using the stubs without probing the runtime package first requires slightly +fewer lookups. + +```toml +[environment] +extra-paths = ["/packages"] +``` + +`/packages/shapes-stubs/polygons/pentagon.pyi`: + +```pyi +class Pentagon: + sides: int + area: float +``` + +`/packages/shapes-stubs/polygons/hexagon.pyi`: + +```pyi +class Hexagon: + sides: int + area: float +``` + +`/packages/shapes/__init__.py`: + +```py +``` + +`/packages/shapes/polygons/__init__.py`: + +```py +``` + +`/packages/shapes/polygons/pentagon.py`: + +```py +class Pentagon: ... +``` + +`/packages/shapes/polygons/hexagon.py`: + +```py +class Hexagon: ... +``` + +`main.py`: + +```py +from shapes.polygons.pentagon import Pentagon +from shapes.polygons.hexagon import Hexagon + +reveal_type(Pentagon().sides) # revealed: int +reveal_type(Hexagon().area) # revealed: int | float +``` + +## Stub package using `__init__.py` over `.pyi` + +It's recommended that stub packages use `__init__.pyi` files over `__init__.py` but it doesn't seem +to be an enforced convention. At least, Pyright is fine with the following. + +```toml +[environment] +extra-paths = ["/packages"] +``` + +`/packages/shapes-stubs/__init__.py`: + +```py +class Pentagon: + sides: int + area: float + +class Hexagon: + sides: int + area: float +``` + +`/packages/shapes/__init__.py`: + +```py +class Pentagon: ... +class Hexagon: ... +``` + +`main.py`: + +```py +from shapes import Hexagon, Pentagon + +reveal_type(Pentagon().sides) # revealed: int +reveal_type(Hexagon().area) # revealed: int | float +``` diff --git a/crates/red_knot_python_semantic/src/module_resolver/module.rs b/crates/red_knot_python_semantic/src/module_resolver/module.rs index 1c5b1192f3575..6256a6f98b12f 100644 --- a/crates/red_knot_python_semantic/src/module_resolver/module.rs +++ b/crates/red_knot_python_semantic/src/module_resolver/module.rs @@ -96,6 +96,9 @@ impl ModuleKind { pub const fn is_package(self) -> bool { matches!(self, ModuleKind::Package) } + pub const fn is_module(self) -> bool { + matches!(self, ModuleKind::Module) + } } /// Enumeration of various core stdlib modules in which important types are located 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 d1ee3869a0db8..6fa67f1a1825d 100644 --- a/crates/red_knot_python_semantic/src/module_resolver/path.rs +++ b/crates/red_knot_python_semantic/src/module_resolver/path.rs @@ -116,8 +116,9 @@ impl ModulePath { | SearchPathInner::SitePackages(search_path) | SearchPathInner::Editable(search_path) => { let absolute_path = search_path.join(relative_path); + system_path_to_file(resolver.db.upcast(), absolute_path.join("__init__.py")).is_ok() - || system_path_to_file(resolver.db.upcast(), absolute_path.join("__init__.py")) + || system_path_to_file(resolver.db.upcast(), absolute_path.join("__init__.pyi")) .is_ok() } SearchPathInner::StandardLibraryCustom(search_path) => { @@ -632,6 +633,19 @@ impl PartialEq for VendoredPathBuf { } } +impl fmt::Display for SearchPath { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match &*self.0 { + SearchPathInner::Extra(system_path_buf) + | SearchPathInner::FirstParty(system_path_buf) + | SearchPathInner::SitePackages(system_path_buf) + | SearchPathInner::Editable(system_path_buf) + | SearchPathInner::StandardLibraryCustom(system_path_buf) => system_path_buf.fmt(f), + SearchPathInner::StandardLibraryVendored(vendored_path_buf) => vendored_path_buf.fmt(f), + } + } +} + #[cfg(test)] mod tests { use ruff_db::Db; 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 31dde32c26c5f..a3ae6d805b3f6 100644 --- a/crates/red_knot_python_semantic/src/module_resolver/resolver.rs +++ b/crates/red_knot_python_semantic/src/module_resolver/resolver.rs @@ -1,6 +1,9 @@ use std::borrow::Cow; +use std::fmt; use std::iter::FusedIterator; +use std::str::Split; +use compact_str::format_compact; use rustc_hash::{FxBuildHasher, FxHashSet}; use ruff_db::files::{File, FilePath, FileRootKind}; @@ -36,16 +39,21 @@ pub(crate) fn resolve_module_query<'db>( let name = module_name.name(db); let _span = tracing::trace_span!("resolve_module", %name).entered(); - let Some((search_path, module_file, kind)) = resolve_name(db, name) else { + let Some((search_path, resolved_module)) = resolve_name(db, name) else { tracing::debug!("Module `{name}` not found in search paths"); return None; }; - let module = Module::new(name.clone(), kind, search_path, module_file); - tracing::trace!( "Resolved module `{name}` to `{path}`", - path = module_file.path(db) + path = resolved_module.file.path(db) + ); + + let module = Module::new( + name.clone(), + resolved_module.kind, + search_path, + resolved_module.file, ); Some(module) @@ -581,13 +589,16 @@ struct ModuleNameIngredient<'db> { /// Given a module name and a list of search paths in which to lookup modules, /// attempt to resolve the module name -fn resolve_name(db: &dyn Db, name: &ModuleName) -> Option<(SearchPath, File, ModuleKind)> { +fn resolve_name(db: &dyn Db, name: &ModuleName) -> Option<(SearchPath, ResolvedModule)> { let program = Program::get(db); let python_version = program.python_version(db); let resolver_state = ResolverContext::new(db, python_version); let is_builtin_module = ruff_python_stdlib::sys::is_builtin_module(python_version.minor, name.as_str()); + let name = RelaxedModuleName::new(name); + let stub_name = name.to_stub_package(); + for search_path in search_paths(db) { // When a builtin module is imported, standard module resolution is bypassed: // the module name always resolves to the stdlib module, @@ -597,45 +608,110 @@ fn resolve_name(db: &dyn Db, name: &ModuleName) -> Option<(SearchPath, File, Mod continue; } - let mut components = name.components(); - let module_name = components.next_back()?; - - match resolve_package(search_path, components, &resolver_state) { - Ok(resolved_package) => { - let mut package_path = resolved_package.path; - - package_path.push(module_name); - - // 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 !search_path.is_standard_library() { + match resolve_module_in_search_path(&resolver_state, &stub_name, search_path) { + Ok(resolved_module) => { + if resolved_module.package_kind.is_root() && resolved_module.kind.is_module() { + tracing::trace!("Search path '{search_path} contains a module named `{stub_name}` but a standalone module isn't a valid stub."); + } else { + return Some((search_path.clone(), resolved_module)); + } } - - // 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)); + Err(PackageKind::Root) => { + tracing::trace!( + "Search path '{search_path}' contains no stub package named `{stub_name}`." + ); } - - // For regular packages, don't search the next search path. All files of that - // package must be in the same location - if resolved_package.kind.is_regular_package() { + Err(PackageKind::Regular) => { + tracing::trace!( + "Stub-package in `{search_path} doesn't contain module: `{name}`" + ); + // stub exists, but the module doesn't. + // TODO: Support partial packages. return None; } + Err(PackageKind::Namespace) => { + tracing::trace!( + "Stub-package in `{search_path} doesn't contain module: `{name}` but it is a namespace package, keep going." + ); + // stub exists, but the module doesn't. But this is a namespace package, + // keep searching the next search path for a stub package with the same name. + continue; + } } - Err(parent_kind) => { - if parent_kind.is_regular_package() { - // For regular packages, don't search the next search path. + } + + match resolve_module_in_search_path(&resolver_state, &name, search_path) { + Ok(resolved_module) => return Some((search_path.clone(), resolved_module)), + Err(kind) => match kind { + PackageKind::Root => { + tracing::trace!( + "Search path '{search_path}' contains no package named `{name}`." + ); + } + PackageKind::Regular => { + // For regular packages, don't search the next search path. All files of that + // package must be in the same location + tracing::trace!("Package in `{search_path} doesn't contain module: `{name}`"); return None; } - } + PackageKind::Namespace => { + tracing::trace!( + "Package in `{search_path} doesn't contain module: `{name}` but it is a namespace package, keep going." + ); + } + }, } } None } +#[derive(Debug)] +struct ResolvedModule { + kind: ModuleKind, + package_kind: PackageKind, + file: File, +} + +fn resolve_module_in_search_path( + context: &ResolverContext, + name: &RelaxedModuleName, + search_path: &SearchPath, +) -> Result { + let mut components = name.components(); + let module_name = components.next_back().unwrap(); + + let resolved_package = resolve_package(search_path, components, context)?; + + let mut package_path = resolved_package.path; + + package_path.push(module_name); + + // Check for a regular package first (highest priority) + package_path.push("__init__"); + if let Some(regular_package) = resolve_file_module(&package_path, context) { + return Ok(ResolvedModule { + file: regular_package, + kind: ModuleKind::Package, + package_kind: resolved_package.kind, + }); + } + + // Check for a file module next + package_path.pop(); + + if let Some(file_module) = resolve_file_module(&package_path, context) { + return Ok(ResolvedModule { + file: file_module, + kind: ModuleKind::Module, + package_kind: resolved_package.kind, + }); + } + + Err(resolved_package.kind) +} + /// If `module` exists on disk with either a `.pyi` or `.py` extension, /// return the [`File`] corresponding to that path. /// @@ -698,7 +774,7 @@ where // 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. + // A directory without an `__init__.py(i)` is a namespace package, continue with the next folder. in_namespace_package = true; } else if in_namespace_package { // Package not found but it is part of a namespace package. @@ -751,8 +827,8 @@ enum PackageKind { } impl PackageKind { - const fn is_regular_package(self) -> bool { - matches!(self, PackageKind::Regular) + pub(crate) const fn is_root(self) -> bool { + matches!(self, PackageKind::Root) } } @@ -771,6 +847,34 @@ impl<'db> ResolverContext<'db> { } } +/// A [`ModuleName`] but with relaxed semantics to allow `-stubs.path` +#[derive(Debug)] +struct RelaxedModuleName(compact_str::CompactString); + +impl RelaxedModuleName { + fn new(name: &ModuleName) -> Self { + Self(name.as_str().into()) + } + + fn components(&self) -> Split<'_, char> { + self.0.split('.') + } + + fn to_stub_package(&self) -> Self { + if let Some((package, rest)) = self.0.split_once('.') { + Self(format_compact!("{package}-stubs.{rest}")) + } else { + Self(format_compact!("{package}-stubs", package = self.0)) + } + } +} + +impl fmt::Display for RelaxedModuleName { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + self.0.fmt(f) + } +} + #[cfg(test)] mod tests { use ruff_db::files::{system_path_to_file, File, FilePath};