Skip to content

Commit d2d9e8a

Browse files
authored
[ty] Preserve namespace fallback for partial stub packages (astral-sh#29559)
Partial stub packages could hide runtime submodules when the corresponding package used `pkgutil.extend_path` or was an implicit namespace package. Preserve runtime namespace candidates while resolving descendants so modules missing from partial stubs can fall back to the implementation. Track concrete stub packages separately from ordinary packages, preserving normal namespace shadowing, complete-stub behavior, and partial subpackages. Closes astral-sh/ty#4767.
1 parent 6b16159 commit d2d9e8a

3 files changed

Lines changed: 240 additions & 33 deletions

File tree

‎crates/ty_module_resolver/src/resolve.rs‎

Lines changed: 50 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -1251,13 +1251,9 @@ enum ComponentFileFilter {
12511251
StubOnly,
12521252
}
12531253

1254-
/// Where a candidate sits in the typing specification's module-resolution order.
1255-
///
1256-
/// Variants are declared from highest to lowest precedence so that derived ordering can be used
1257-
/// when traversing candidates. This is a precedence tier rather than a total ordering: the stable
1258-
/// sorts used by the resolver preserve search-path order between candidates in the same tier.
1259-
#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, get_size2::GetSize)]
1260-
enum CandidatePrecedence {
1254+
/// Where a module resolution candidate originated.
1255+
#[derive(Debug, Copy, Clone, PartialEq, Eq, get_size2::GetSize)]
1256+
enum CandidateOrigin {
12611257
/// A PEP 561 stub-only package named `<package>-stubs`.
12621258
///
12631259
/// Stub packages take precedence over candidates for `<package>` regardless of where those
@@ -1268,7 +1264,19 @@ enum CandidatePrecedence {
12681264
///
12691265
/// This includes `.pyi` and `.py` packages and modules from extra paths, first-party code,
12701266
/// editable installs, site-packages, and the standard library.
1271-
SearchPathOrder,
1267+
SearchPath,
1268+
}
1269+
1270+
impl CandidateOrigin {
1271+
/// Returns the candidate's precedence tier in the typing specification's module-resolution order.
1272+
///
1273+
/// The resolver's stable sort preserves search-path order between candidates in the same tier.
1274+
fn sort_key(self) -> impl Ord {
1275+
match self {
1276+
Self::StubPackage => 0,
1277+
Self::SearchPath => 1,
1278+
}
1279+
}
12721280
}
12731281

12741282
#[derive(Debug, Clone)]
@@ -1280,28 +1288,28 @@ struct ModuleResolutionCandidate<'db> {
12801288
directory: ModuleDirectory<'db>,
12811289
module: ResolvedModule,
12821290
py_typed: PyTyped,
1283-
precedence: CandidatePrecedence,
1291+
origin: CandidateOrigin,
12841292
}
12851293

12861294
impl<'db> ModuleResolutionCandidate<'db> {
12871295
fn root(context: &ResolverContext<'db>, search_path: &SearchPath) -> Self {
1288-
Self::with_precedence(context, search_path, CandidatePrecedence::SearchPathOrder)
1296+
Self::with_origin(context, search_path, CandidateOrigin::SearchPath)
12891297
}
12901298

12911299
fn stub(context: &ResolverContext<'db>, search_path: &SearchPath) -> Self {
1292-
Self::with_precedence(context, search_path, CandidatePrecedence::StubPackage)
1300+
Self::with_origin(context, search_path, CandidateOrigin::StubPackage)
12931301
}
12941302

1295-
fn with_precedence(
1303+
fn with_origin(
12961304
context: &ResolverContext<'db>,
12971305
search_path: &SearchPath,
1298-
precedence: CandidatePrecedence,
1306+
origin: CandidateOrigin,
12991307
) -> Self {
13001308
Self {
13011309
directory: ModuleDirectory::new(context, search_path.to_module_path(), Some(true)),
13021310
module: ResolvedModule::NamespacePackage,
13031311
py_typed: PyTyped::Untyped,
1304-
precedence,
1312+
origin,
13051313
}
13061314
}
13071315

@@ -1414,9 +1422,11 @@ fn resolve_stub_package_in_search_path<'db>(
14141422
/// Orders candidates for the same module name by precedence and removes shadowed namespace packages.
14151423
///
14161424
/// Regular packages and file modules shadow namespace packages, including legacy namespaces.
1417-
/// When `for_module_name_prefix` is true, partial namespaces with higher precedence than every
1418-
/// competing regular package or file module remain available to supply descendants. Candidates
1419-
/// within the same precedence tier retain their search-path order.
1425+
/// While resolving descendants, stub packages retain ordinary namespace candidates as fallbacks,
1426+
/// and partial namespace stubs remain available alongside ordinary concrete candidates. Within
1427+
/// each precedence tier, concrete candidates still shadow namespace packages. At the final
1428+
/// component, any concrete candidate shadows namespaces. Candidates within the same precedence
1429+
/// tier retain their search-path order.
14201430
fn normalize_candidates<'db>(
14211431
context: &ResolverContext,
14221432
mut candidates: ResolvedNames<'db>,
@@ -1427,13 +1437,19 @@ fn normalize_candidates<'db>(
14271437
return candidates;
14281438
}
14291439

1430-
let best_concrete_precedence = candidates
1440+
let mut has_concrete_stub_package = false;
1441+
let mut has_concrete_search_path_candidate = false;
1442+
for candidate in candidates
14311443
.iter()
14321444
.filter(|candidate| !candidate.is_any_namespace_package(context))
1433-
.map(|candidate| candidate.precedence)
1434-
.min();
1445+
{
1446+
match candidate.origin {
1447+
CandidateOrigin::StubPackage => has_concrete_stub_package = true,
1448+
CandidateOrigin::SearchPath => has_concrete_search_path_candidate = true,
1449+
}
1450+
}
14351451

1436-
candidates.sort_by_key(|candidate| candidate.precedence);
1452+
candidates.sort_by_key(|candidate| candidate.origin.sort_key());
14371453

14381454
// Note that we intentionally do *not* filter out ordinary search-path candidates when a stub
14391455
// package is found. Even when a non-namespace, non-partial stub package exists, we keep the
@@ -1445,15 +1461,20 @@ fn normalize_candidates<'db>(
14451461
return true;
14461462
}
14471463

1448-
// A higher-precedence partial namespace remains available while resolving its descendants.
1449-
// At the final component, a concrete package or module shadows it.
1450-
let preserved_for_descendants = best_concrete_precedence.is_none_or(|precedence| {
1451-
for_module_name_prefix
1452-
&& candidate.py_typed == PyTyped::Partial
1453-
&& candidate.precedence < precedence
1454-
});
1464+
let shadowed = if for_module_name_prefix {
1465+
match candidate.origin {
1466+
CandidateOrigin::StubPackage => {
1467+
has_concrete_stub_package
1468+
|| (has_concrete_search_path_candidate
1469+
&& candidate.py_typed != PyTyped::Partial)
1470+
}
1471+
CandidateOrigin::SearchPath => has_concrete_search_path_candidate,
1472+
}
1473+
} else {
1474+
has_concrete_stub_package || has_concrete_search_path_candidate
1475+
};
14551476

1456-
if preserved_for_descendants {
1477+
if !shadowed {
14571478
return true;
14581479
}
14591480

‎crates/ty_module_resolver/src/resolve/search.rs‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ use crate::module_name::ModuleName;
3030
use crate::path::{ModuleDirectory, ModuleDirectoryEntry, ModulePath, SearchPath};
3131

3232
use super::{
33-
CandidatePrecedence, ComponentFileFilter, ModuleNameIngredient, ModuleResolutionCandidate,
33+
CandidateOrigin, ComponentFileFilter, ModuleNameIngredient, ModuleResolutionCandidate,
3434
ModuleResolveMode, PyTyped, ResolvedModule, ResolvedNames, ResolverContext, StubPackageIndex,
3535
StubPackagePaths, normalize_candidates, resolve_component, resolve_stub_package_in_search_path,
3636
search_paths, stub_package_index,
@@ -509,7 +509,7 @@ struct CachedCandidate {
509509
path: ModulePath,
510510
module: ResolvedModule,
511511
py_typed: PyTyped,
512-
precedence: CandidatePrecedence,
512+
origin: CandidateOrigin,
513513

514514
/// Stores the result of the symlink policy for enumeration. This prevents
515515
/// unnecessary cache invalidation of a `ModuleListing`, since the symlink
@@ -529,7 +529,7 @@ impl CachedCandidate {
529529
enumeration_allowed: candidate.directory.enumeration_allowed(db),
530530
module: candidate.module,
531531
py_typed: candidate.py_typed,
532-
precedence: candidate.precedence,
532+
origin: candidate.origin,
533533
}
534534
}
535535

@@ -542,7 +542,7 @@ impl CachedCandidate {
542542
),
543543
module: self.module,
544544
py_typed: self.py_typed,
545-
precedence: self.precedence,
545+
origin: self.origin,
546546
}
547547
}
548548
}

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

Lines changed: 186 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -469,3 +469,189 @@ reveal_type(Both().both) # revealed: str
469469
reveal_type(Impl().impl) # revealed: str
470470
reveal_type(Fake().fake) # revealed: Unknown
471471
```
472+
473+
## Partial stub for a legacy namespace package
474+
475+
A partial stub package takes precedence over a legacy namespace package for modules it defines.
476+
Subpackages missing from the stubs still resolve from the runtime package.
477+
478+
Regression test for <https://github.com/astral-sh/ty/issues/4767>.
479+
480+
```toml
481+
[environment]
482+
extra-paths = ["/packages"]
483+
```
484+
485+
`/packages/foo-stubs/py.typed`:
486+
487+
```text
488+
partial
489+
```
490+
491+
`/packages/foo-stubs/__init__.pyi`:
492+
493+
```pyi
494+
x: int
495+
```
496+
497+
`/packages/foo/__init__.py`:
498+
499+
```py
500+
__path__ = __import__("pkgutil").extend_path(__path__, __name__)
501+
x = "runtime"
502+
```
503+
504+
`/packages/foo/sub/__init__.py`:
505+
506+
```py
507+
y = 1
508+
```
509+
510+
`main.py`:
511+
512+
```py
513+
from foo import x
514+
from foo.sub import y
515+
516+
reveal_type(x) # revealed: int
517+
reveal_type(y) # revealed: Literal[1]
518+
```
519+
520+
## Partial stub for an implicit namespace package
521+
522+
The runtime package can also be an implicit namespace package, without an `__init__.py`.
523+
524+
```toml
525+
[environment]
526+
extra-paths = ["/packages"]
527+
```
528+
529+
`/packages/foo-stubs/py.typed`:
530+
531+
```text
532+
partial
533+
```
534+
535+
`/packages/foo-stubs/__init__.pyi`:
536+
537+
```pyi
538+
```
539+
540+
`/packages/foo/impl.py`:
541+
542+
```py
543+
x = 1
544+
```
545+
546+
`main.py`:
547+
548+
```py
549+
from foo.impl import x
550+
551+
reveal_type(x) # revealed: Literal[1]
552+
```
553+
554+
## Complete stub with a partial subpackage in a namespace package
555+
556+
Although the root stub package is complete, its `bar` subpackage declares itself partial. Missing
557+
modules within `bar` still resolve from the runtime package, even when its root is a namespace
558+
package. Runtime-only siblings of `bar` remain inaccessible through the complete root stub.
559+
560+
```toml
561+
[environment]
562+
extra-paths = ["/packages"]
563+
```
564+
565+
`/packages/foo-stubs/__init__.pyi`:
566+
567+
```pyi
568+
```
569+
570+
`/packages/foo-stubs/bar/py.typed`:
571+
572+
```text
573+
partial
574+
```
575+
576+
`/packages/foo-stubs/bar/__init__.pyi`:
577+
578+
```pyi
579+
```
580+
581+
`/packages/foo/__init__.py`:
582+
583+
```py
584+
__path__ = __import__("pkgutil").extend_path(__path__, __name__)
585+
```
586+
587+
`/packages/foo/bar/__init__.py`:
588+
589+
```py
590+
```
591+
592+
`/packages/foo/bar/impl.py`:
593+
594+
```py
595+
x = 1
596+
```
597+
598+
`/packages/foo/hidden.py`:
599+
600+
```py
601+
```
602+
603+
`main.py`:
604+
605+
```py
606+
import foo.hidden # error: [unresolved-import]
607+
from foo.bar.impl import x
608+
609+
reveal_type(x) # revealed: Literal[1]
610+
```
611+
612+
## Partial stubs preserve runtime package precedence
613+
614+
A regular runtime package shadows a runtime namespace package, even when the namespace comes first
615+
on the search path. Partial stubs fall back to the regular package without exposing modules from the
616+
shadowed namespace.
617+
618+
```toml
619+
[environment]
620+
extra-paths = ["/first", "/second"]
621+
```
622+
623+
`/first/foo-stubs/py.typed`:
624+
625+
```text
626+
partial
627+
```
628+
629+
`/first/foo-stubs/__init__.pyi`:
630+
631+
```pyi
632+
```
633+
634+
`/first/foo/hidden.py`:
635+
636+
```py
637+
```
638+
639+
`/second/foo/__init__.py`:
640+
641+
```py
642+
```
643+
644+
`/second/foo/impl.py`:
645+
646+
```py
647+
x = 1
648+
```
649+
650+
`main.py`:
651+
652+
```py
653+
import foo.hidden # error: [unresolved-import]
654+
from foo.impl import x
655+
656+
reveal_type(x) # revealed: Literal[1]
657+
```

0 commit comments

Comments
 (0)