[ty] Implement partial stubs (#19931)

Fixes https://github.com/astral-sh/ty/issues/184

crates/ty_python_semantic/resources/mdtest/import/partial_stub_packages.md

@@ -0,0 +1,425 @@
+# Partial stub packages
+
+Partial stub packages are stubs that are allowed to be missing modules. See the
+[specification](https://peps.python.org/pep-0561/#partial-stub-packages). Partial stubs are also
+called "incomplete" packages, and non-partial stubs are called "complete" packages.
+
+Normally a stub package is expected to define a copy of every module the real implementation
+defines. Module resolution is consequently required to report a module doesn't exist if it finds
+`mypackage-stubs` and fails to find `mypackage.mymodule` *even if* `mypackage` does define
+`mymodule`.
+
+If a stub package declares that it's partial, we instead are expected to fall through to the
+implementation package and try to discover `mymodule` there. This is described as follows:
+
+> Type checkers should merge the stub package and runtime package or typeshed directories. This can
+> be thought of as the functional equivalent of copying the stub package into the same directory as
+> the corresponding runtime package or typeshed folder and type checking the combined directory
+> structure. Thus type checkers MUST maintain the normal resolution order of checking `*.pyi` before
+> `*.py` files.
+
+Namespace stub packages are always considered partial by necessity. Regular stub packages are only
+considered partial if they define a `py.typed` file containing the string `partial\n` (due to real
+stubs in the wild, we relax this and look case-insensitively for `partial`).
+
+The `py.typed` file was originally specified as an empty marker for "this package supports types",
+as a way to opt into having typecheckers run on a package. However ty and pyright choose to largely
+ignore this and just type check every package.
+
+In its original specification it was specified that subpackages inherit any `py.typed` declared in a
+parent package. However the precise interaction with `partial\n` was never specified. We currently
+implement a simple inheritance scheme where a subpackage can always declare its own `py.typed` and
+override whether it's partial or not.
+
+## Partial stub with missing module
+
+A stub package that includes a partial `py.typed` file.
+
+Here "both" is found in the stub, while "impl" is found in the implementation. "fake" is found in
+neither.
+
+```toml
+[environment]
+extra-paths = ["/packages"]
+```
+
+`/packages/foo-stubs/py.typed`:
+
+```text
+partial
+```
+
+`/packages/foo-stubs/__init__.pyi`:
+
+```pyi
+```
+
+`/packages/foo-stubs/both.py`:
+
+```py
+class Both:
+    both: str
+    other: int
+```
+
+`/packages/foo/__init__.py`:
+
+```py
+```
+
+`/packages/foo/both.py`:
+
+```py
+class Both: ...
+```
+
+`/packages/foo/impl.py`:
+
+```py
+class Impl:
+    impl: str
+    other: int
+```
+
+`main.py`:
+
+```py
+from foo.both import Both
+from foo.impl import Impl
+from foo.fake import Fake  # error: "Cannot resolve"
+
+reveal_type(Both().both)  # revealed: str
+reveal_type(Impl().impl)  # revealed: str
+reveal_type(Fake().fake)  # revealed: Unknown
+```
+
+## Non-partial stub with missing module
+
+Omitting the partial `py.typed`, we see "impl" now cannot be found.
+
+```toml
+[environment]
+extra-paths = ["/packages"]
+```
+
+`/packages/foo-stubs/__init__.pyi`:
+
+```pyi
+```
+
+`/packages/foo-stubs/both.py`:
+
+```py
+class Both:
+    both: str
+    other: int
+```
+
+`/packages/foo/__init__.py`:
+
+```py
+```
+
+`/packages/foo/both.py`:
+
+```py
+class Both: ...
+```
+
+`/packages/foo/impl.py`:
+
+```py
+class Impl:
+    impl: str
+    other: int
+```
+
+`main.py`:
+
+```py
+from foo.both import Both
+from foo.impl import Impl  # error: "Cannot resolve"
+from foo.fake import Fake  # error: "Cannot resolve"
+
+reveal_type(Both().both)  # revealed: str
+reveal_type(Impl().impl)  # revealed: Unknown
+reveal_type(Fake().fake)  # revealed: Unknown
+```
+
+## Full-typed stub with missing module
+
+Including a blank py.typed we still don't conclude it's partial.
+
+```toml
+[environment]
+extra-paths = ["/packages"]
+```
+
+`/packages/foo-stubs/py.typed`:
+
+```text
+```
+
+`/packages/foo-stubs/__init__.pyi`:
+
+```pyi
+```
+
+`/packages/foo-stubs/both.py`:
+
+```py
+class Both:
+    both: str
+    other: int
+```
+
+`/packages/foo/__init__.py`:
+
+```py
+```
+
+`/packages/foo/both.py`:
+
+```py
+class Both: ...
+```
+
+`/packages/foo/impl.py`:
+
+```py
+class Impl:
+    impl: str
+    other: int
+```
+
+`main.py`:
+
+```py
+from foo.both import Both
+from foo.impl import Impl  # error: "Cannot resolve"
+from foo.fake import Fake  # error: "Cannot resolve"
+
+reveal_type(Both().both)  # revealed: str
+reveal_type(Impl().impl)  # revealed: Unknown
+reveal_type(Fake().fake)  # revealed: Unknown
+```
+
+## Inheriting a partial `py.typed`
+
+`foo-stubs` defines a partial py.typed which is used by `foo-stubs/bar`
+
+```toml
+[environment]
+extra-paths = ["/packages"]
+```
+
+`/packages/foo-stubs/py.typed`:
+
+```text
+# Also testing permissive parsing
+# PARTIAL\n
+```
+
+`/packages/foo-stubs/__init__.pyi`:
+
+```pyi
+```
+
+`/packages/foo-stubs/bar/__init__.pyi`:
+
+```pyi
+```
+
+`/packages/foo-stubs/bar/both.py`:
+
+```py
+class Both:
+    both: str
+    other: int
+```
+
+`/packages/foo/__init__.py`:
+
+```py
+```
+
+`/packages/foo/bar/__init__.py`:
+
+```py
+```
+
+`/packages/foo/bar/both.py`:
+
+```py
+class Both: ...
+```
+
+`/packages/foo/bar/impl.py`:
+
+```py
+class Impl:
+    impl: str
+    other: int
+```
+
+`main.py`:
+
+```py
+from foo.bar.both import Both
+from foo.bar.impl import Impl
+from foo.bar.fake import Fake  # error: "Cannot resolve"
+
+reveal_type(Both().both)  # revealed: str
+reveal_type(Impl().impl)  # revealed: str
+reveal_type(Fake().fake)  # revealed: Unknown
+```
+
+## Overloading a full `py.typed`
+
+`foo-stubs` defines a full py.typed which is overloaded to partial by `foo-stubs/bar`
+
+```toml
+[environment]
+extra-paths = ["/packages"]
+```
+
+`/packages/foo-stubs/py.typed`:
+
+```text
+```
+
+`/packages/foo-stubs/bar/py.typed`:
+
+```text
+# Also testing permissive parsing
+partial/n
+```
+
+`/packages/foo-stubs/__init__.pyi`:
+
+```pyi
+```
+
+`/packages/foo-stubs/bar/__init__.pyi`:
+
+```pyi
+```
+
+`/packages/foo-stubs/bar/both.py`:
+
+```py
+class Both:
+    both: str
+    other: int
+```
+
+`/packages/foo/__init__.py`:
+
+```py
+```
+
+`/packages/foo/bar/__init__.py`:
+
+```py
+```
+
+`/packages/foo/bar/both.py`:
+
+```py
+class Both: ...
+```
+
+`/packages/foo/bar/impl.py`:
+
+```py
+class Impl:
+    impl: str
+    other: int
+```
+
+`main.py`:
+
+```py
+from foo.bar.both import Both
+from foo.bar.impl import Impl
+from foo.bar.fake import Fake  # error: "Cannot resolve"
+
+reveal_type(Both().both)  # revealed: str
+reveal_type(Impl().impl)  # revealed: str
+reveal_type(Fake().fake)  # revealed: Unknown
+```
+
+## Overloading a partial `py.typed`
+
+`foo-stubs` defines a partial py.typed which is overloaded to full by `foo-stubs/bar`
+
+```toml
+[environment]
+extra-paths = ["/packages"]
+```
+
+`/packages/foo-stubs/py.typed`:
+
+```text
+# Also testing permissive parsing
+pArTiAl\n
+```
+
+`/packages/foo-stubs/bar/py.typed`:
+
+```text
+```
+
+`/packages/foo-stubs/__init__.pyi`:
+
+```pyi
+```
+
+`/packages/foo-stubs/bar/__init__.pyi`:
+
+```pyi
+```
+
+`/packages/foo-stubs/bar/both.py`:
+
+```py
+class Both:
+    both: str
+    other: int
+```
+
+`/packages/foo/__init__.py`:
+
+```py
+```
+
+`/packages/foo/bar/__init__.py`:
+
+```py
+```
+
+`/packages/foo/bar/both.py`:
+
+```py
+class Both: ...
+```
+
+`/packages/foo/bar/impl.py`:
+
+```py
+class Impl:
+    impl: str
+    other: int
+```
+
+`main.py`:
+
+```py
+from foo.bar.both import Both
+from foo.bar.impl import Impl  # error: "Cannot resolve"
+from foo.bar.fake import Fake  # error: "Cannot resolve"
+
+reveal_type(Both().both)  # revealed: str
+reveal_type(Impl().impl)  # revealed: Unknown
+reveal_type(Fake().fake)  # revealed: Unknown
+```

crates/ty_python_semantic/src/module_resolver/path.rs

@@ -11,7 +11,7 @@ use ruff_db::vendored::{VendoredPath, VendoredPathBuf};
 use super::typeshed::{TypeshedVersionsParseError, TypeshedVersionsQueryResult, typeshed_versions};
 use crate::Db;
 use crate::module_name::ModuleName;
-use crate::module_resolver::resolver::ResolverContext;
+use crate::module_resolver::resolver::{PyTyped, ResolverContext};
 use crate::site_packages::SitePackagesDiscoveryError;
 
 /// A path that points to a Python module.
@@ -148,6 +148,31 @@ impl ModulePath {
         }
     }
 
+    /// Get the `py.typed` info for this package (not considering parent packages)
+    pub(super) fn py_typed(&self, resolver: &ResolverContext) -> PyTyped {
+        let Some(py_typed_contents) = self.to_system_path().and_then(|path| {
+            let py_typed_path = path.join("py.typed");
+            let py_typed_file = system_path_to_file(resolver.db, py_typed_path).ok()?;
+            // If we fail to read it let's say that's like it doesn't exist
+            // (right now the difference between Untyped and Full is academic)
+            py_typed_file.read_to_string(resolver.db).ok()
+        }) else {
+            return PyTyped::Untyped;
+        };
+        // The python typing spec says to look for "partial\n" but in the wild we've seen:
+        //
+        // * PARTIAL\n
+        // * partial\\n (as in they typed "\n")
+        // * partial/n
+        //
+        // since the py.typed file never really grew any other contents, let's be permissive
+        if py_typed_contents.to_ascii_lowercase().contains("partial") {
+            PyTyped::Partial
+        } else {
+            PyTyped::Full
+        }
+    }
+
     pub(super) fn to_system_path(&self) -> Option<SystemPathBuf> {
         let ModulePath {
             search_path,

crates/ty_python_semantic/src/module_resolver/resolver.rs

@@ -684,7 +684,7 @@ fn resolve_name(db: &dyn Db, name: &ModuleName, mode: ModuleResolveMode) -> Opti
 
         if !search_path.is_standard_library() && resolver_state.mode.stubs_allowed() {
             match resolve_name_in_search_path(&resolver_state, &stub_name, search_path) {
-                Ok((package_kind, ResolvedName::FileModule(module))) => {
+                Ok((package_kind, _, ResolvedName::FileModule(module))) => {
                     if package_kind.is_root() && module.kind.is_module() {
                         tracing::trace!(
                             "Search path `{search_path}` contains a module \
@@ -694,23 +694,30 @@ fn resolve_name(db: &dyn Db, name: &ModuleName, mode: ModuleResolveMode) -> Opti
                         return Some(ResolvedName::FileModule(module));
                     }
                 }
-                Ok((_, ResolvedName::NamespacePackage)) => {
+                Ok((_, _, ResolvedName::NamespacePackage)) => {
                     is_namespace_package = true;
                 }
-                Err(PackageKind::Root) => {
+                Err((PackageKind::Root, _)) => {
                     tracing::trace!(
                         "Search path `{search_path}` contains no stub package named `{stub_name}`."
                     );
                 }
-                Err(PackageKind::Regular) => {
+                Err((PackageKind::Regular, PyTyped::Partial)) => {
+                    tracing::trace!(
+                        "Stub-package in `{search_path}` doesn't contain module: \
+                         `{name}` but it is a partial package, keep going."
+                    );
+                    // stub exists, but the module doesn't. But this is a partial package,
+                    // fall through to looking for a non-stub 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) => {
+                Err((PackageKind::Namespace, _)) => {
                     tracing::trace!(
                         "Stub-package in `{search_path}` doesn't contain module: \
                          `{name}` but it is a namespace package, keep going."
@@ -723,25 +730,31 @@ fn resolve_name(db: &dyn Db, name: &ModuleName, mode: ModuleResolveMode) -> Opti
         }
 
         match resolve_name_in_search_path(&resolver_state, &name, search_path) {
-            Ok((_, ResolvedName::FileModule(module))) => {
+            Ok((_, _, ResolvedName::FileModule(module))) => {
                 return Some(ResolvedName::FileModule(module));
             }
-            Ok((_, ResolvedName::NamespacePackage)) => {
+            Ok((_, _, ResolvedName::NamespacePackage)) => {
                 is_namespace_package = true;
             }
             Err(kind) => match kind {
-                PackageKind::Root => {
+                (PackageKind::Root, _) => {
                     tracing::trace!(
                         "Search path `{search_path}` contains no package named `{name}`."
                     );
                 }
-                PackageKind::Regular => {
+                (PackageKind::Regular, PyTyped::Partial) => {
+                    tracing::trace!(
+                        "Package in `{search_path}` doesn't contain module: \
+                         `{name}` but it is a partial package, keep going."
+                    );
+                }
+                (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 => {
+                (PackageKind::Namespace, _) => {
                     tracing::trace!(
                         "Package in `{search_path}` doesn't contain module: \
                          `{name}` but it is a namespace package, keep going."
@@ -796,7 +809,7 @@ fn resolve_name_in_search_path(
     context: &ResolverContext,
     name: &RelaxedModuleName,
     search_path: &SearchPath,
-) -> Result<(PackageKind, ResolvedName), PackageKind> {
+) -> Result<(PackageKind, PyTyped, ResolvedName), (PackageKind, PyTyped)> {
     let mut components = name.components();
     let module_name = components.next_back().unwrap();
 
@@ -811,6 +824,7 @@ fn resolve_name_in_search_path(
     if let Some(regular_package) = resolve_file_module(&package_path, context) {
         return Ok((
             resolved_package.kind,
+            resolved_package.typed,
             ResolvedName::FileModule(ResolvedFileModule {
                 search_path: search_path.clone(),
                 kind: ModuleKind::Package,
@@ -825,6 +839,7 @@ fn resolve_name_in_search_path(
     if let Some(file_module) = resolve_file_module(&package_path, context) {
         return Ok((
             resolved_package.kind,
+            resolved_package.typed,
             ResolvedName::FileModule(ResolvedFileModule {
                 file: file_module,
                 kind: ModuleKind::Module,
@@ -859,12 +874,16 @@ fn resolve_name_in_search_path(
                     package_path.search_path().as_system_path().unwrap(),
                 )
             {
-                return Ok((resolved_package.kind, ResolvedName::NamespacePackage));
+                return Ok((
+                    resolved_package.kind,
+                    resolved_package.typed,
+                    ResolvedName::NamespacePackage,
+                ));
             }
         }
     }
 
-    Err(resolved_package.kind)
+    Err((resolved_package.kind, resolved_package.typed))
 }
 
 /// If `module` exists on disk with either a `.pyi` or `.py` extension,
@@ -919,7 +938,7 @@ fn resolve_package<'a, 'db, I>(
     module_search_path: &SearchPath,
     components: I,
     resolver_state: &ResolverContext<'db>,
-) -> Result<ResolvedPackage, PackageKind>
+) -> Result<ResolvedPackage, (PackageKind, PyTyped)>
 where
     I: Iterator<Item = &'a str>,
 {
@@ -933,9 +952,12 @@ where
     // `true` if resolving a sub-package. For example, `true` when resolving `bar` of `foo.bar`.
     let mut in_sub_package = false;
 
+    let mut typed = package_path.py_typed(resolver_state);
+
     // For `foo.bar.baz`, test that `foo` and `bar` both contain a `__init__.py`.
     for folder in components {
         package_path.push(folder);
+        typed = package_path.py_typed(resolver_state).inherit_parent(typed);
 
         let is_regular_package = package_path.is_regular_package(resolver_state);
 
@@ -950,13 +972,13 @@ where
             in_namespace_package = true;
         } else if in_namespace_package {
             // Package not found but it is part of a namespace package.
-            return Err(PackageKind::Namespace);
+            return Err((PackageKind::Namespace, typed));
         } else if in_sub_package {
             // A regular sub package wasn't found.
-            return Err(PackageKind::Regular);
+            return Err((PackageKind::Regular, typed));
         } else {
             // We couldn't find `foo` for `foo.bar.baz`, search the next search path.
-            return Err(PackageKind::Root);
+            return Err((PackageKind::Root, typed));
         }
 
         in_sub_package = true;
@@ -973,6 +995,7 @@ where
     Ok(ResolvedPackage {
         kind,
         path: package_path,
+        typed,
     })
 }
 
@@ -980,6 +1003,7 @@ where
 struct ResolvedPackage {
     path: ModulePath,
     kind: PackageKind,
+    typed: PyTyped,
 }
 
 #[derive(Copy, Clone, Eq, PartialEq, Debug)]
@@ -1006,6 +1030,32 @@ impl PackageKind {
     }
 }
 
+/// Info about the `py.typed` file for this package
+#[derive(Copy, Clone, Eq, PartialEq, Debug)]
+pub(crate) enum PyTyped {
+    /// No `py.typed` was found
+    Untyped,
+    /// A `py.typed` was found containing "partial"
+    Partial,
+    /// A `py.typed` was found (not partial)
+    Full,
+}
+
+impl PyTyped {
+    /// Inherit py.typed info from the parent package
+    ///
+    /// > This marker applies recursively: if a top-level package includes it,
+    /// > all its sub-packages MUST support type checking as well.
+    ///
+    /// This implementation implies that once a `py.typed` is specified
+    /// all child packages inherit it, so they can never become Untyped.
+    /// However they can override whether that's Full or Partial by
+    /// redeclaring a `py.typed` file of their own.
+    fn inherit_parent(self, parent: Self) -> Self {
+        if self == Self::Untyped { parent } else { self }
+    }
+}
+
 pub(super) struct ResolverContext<'db> {
     pub(super) db: &'db dyn Db,
     pub(super) python_version: PythonVersion,