mirrors/ruff
1
0
Fork 0
mirror of https://github.com/astral-sh/ruff.git synced 2025-08-04 10:48:32 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

[red-knot] Support stub packages (#17204)

Browse source 
## Summary

This PR adds support for stub packages, except for partial stub packages
(a stub package is always considered non-partial).

I read the specification at
[typing.python.org/en/latest/spec/distributing.html#stub-only-packages](https://typing.python.org/en/latest/spec/distributing.html#stub-only-packages)
but I found it lacking some details, especially on how to handle
namespace packages or when the regular and stub packages disagree on
whether they're namespace packages. I tried to document my decisions in
the mdtests where the specification isn't clear and compared the
behavior to Pyright.

Mypy seems to only support stub packages in the venv folder. At least,
it never picked up my stub packages otherwise. I decided not to spend
too much time fighting mypyp, which is why I focused the comparison
around Pyright

Closes https://github.com/astral-sh/ruff/issues/16612

## Test plan

Added mdtests
This commit is contained in:
Micha Reiser 2025-04-07 14:40:50 +02:00 • committed by GitHub
parent c12c76e9c8
commit 6cc2d02dfa
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 443 additions and 36 deletions
Download patch file Download diff file Expand all files Collapse all files

286
crates/red_knot_python_semantic/resources/mdtest/import/stub_packages.md Normal file
View file

@ -0,0 +1,286 @@
# Stub packages
Stub packages are packages named `<package>-stubs` that provide typing stubs for `<package>`. 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 `<module>-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
```

3
crates/red_knot_python_semantic/src/module_resolver/module.rs
View file

@ -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

16
crates/red_knot_python_semantic/src/module_resolver/path.rs
View file

@ -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<SearchPath> 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;

174
crates/red_knot_python_semantic/src/module_resolver/resolver.rs
View file

@ -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(parent_kind) => {
if parent_kind.is_regular_package() {
// For regular packages, don't search the next search path.
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;
}
}
}
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<ResolvedModule, PackageKind> {
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 `<package>-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};