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

## 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 #16612

## Test plan

Added mdtests

Author: MichaReiser

crates/red_knot_python_semantic/resources/mdtest/import/stub_packages.md

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

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

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