[ty] introduce local variables for from imports of submodules in __init__.py(i) (#21173)

This rips out the previous implementation in favour of a new
implementation with 3 rules:

- **froms are locals**: a `from..import` can only define locals, it does
not have global
side-effects. Specifically any submodule attribute `a` that's implicitly
introduced by either
`from .a import b` or `from . import a as b` (in an `__init__.py(i)`) is
a local and not a
global. If you do such an import at the top of a file you won't notice
this. However if you do
such an import in a function, that means it will only be function-scoped
(so you'll need to do
it in every function that wants to access it, making your code less
sensitive to execution
    order).

- **first from first serve**: only the *first* `from..import` in an
`__init__.py(i)` that imports a
particular direct submodule of the current package introduces that
submodule as a local.
Subsequent imports of the submodule will not introduce that local. This
reflects the fact that
in actual python only the first import of a submodule (in the entire
execution of the program)
introduces it as an attribute of the package. By "first" we mean "the
first time in this scope
(or any parent scope)". This pairs well with the fact that we are
specifically introducing a
local (as long as you don't accidentally shadow or overwrite the local).

- **dot re-exports**: `from . import a` in an `__init__.pyi` is
considered a re-export of `a`
(equivalent to `from . import a as a`). This is required to properly
handle many stubs in the
    wild. Currently it must be *exactly* `from . import ...`.
    
This implementation is intentionally limited/conservative (notably,
often requiring a from import to be relative). I'm going to file a ton
of followups for improvements so that their impact can be evaluated
separately.


Fixes https://github.com/astral-sh/ty/issues/133
This commit is contained in:
Aria Desires 2025-11-10 18:04:56 -05:00 committed by GitHub
parent 1fd852fb3f
commit 2bc6c78e26
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 658 additions and 193 deletions

View file

@ -475,6 +475,12 @@ impl File {
self.path(db).as_str().ends_with("__init__.pyi")
}
/// Returns `true` if the file is an `__init__.pyi`
pub fn is_package(self, db: &dyn Db) -> bool {
let path = self.path(db).as_str();
path.ends_with("__init__.pyi") || path.ends_with("__init__.py")
}
pub fn source_type(self, db: &dyn Db) -> PySourceType {
match self.path(db) {
FilePath::System(path) => path

View file

@ -1,39 +1,39 @@
# Nonstandard Import Conventions
This document covers ty-specific extensions to the
[standard import conventions](https://typing.python.org/en/latest/spec/distributing.html#import-conventions).
[standard import conventions](https://typing.python.org/en/latest/spec/distributing.html#import-conventions),
and other intentional deviations from actual python semantics.
It's a common idiom for a package's `__init__.py(i)` to include several imports like
`from . import mysubmodule`, with the intent that the `mypackage.mysubmodule` attribute should work
for anyone who only imports `mypackage`.
This file currently covers the following details:
In the context of a `.py` we handle this well through our general attempts to faithfully implement
import side-effects. However for `.pyi` files we are expected to apply
[a more strict set of rules](https://typing.python.org/en/latest/spec/distributing.html#import-conventions)
to encourage intentional API design. Although `.pyi` files are explicitly designed to work with
typecheckers, which ostensibly should all enforce these strict rules, every typechecker has its own
defacto "extensions" to them and so a few idioms like `from . import mysubmodule` have found their
way into `.pyi` files too.
- **froms are locals**: a `from..import` can only define locals, it does not have global
side-effects. Specifically any submodule attribute `a` that's implicitly introduced by either
`from .a import b` or `from . import a as b` (in an `__init__.py(i)`) is a local and not a
global. If you do such an import at the top of a file you won't notice this. However if you do
such an import in a function, that means it will only be function-scoped (so you'll need to do
it in every function that wants to access it, making your code less sensitive to execution
order).
Thus for the sake of compatibility, we need to define our own "extensions". Any extensions we define
here have several competing concerns:
- **first from first serve**: only the *first* `from..import` in an `__init__.py(i)` that imports a
particular direct submodule of the current package introduces that submodule as a local.
Subsequent imports of the submodule will not introduce that local. This reflects the fact that
in actual python only the first import of a submodule (in the entire execution of the program)
introduces it as an attribute of the package. By "first" we mean "the first time in this scope
(or any parent scope)". This pairs well with the fact that we are specifically introducing a
local (as long as you don't accidentally shadow or overwrite the local).
- Extensions should ideally be kept narrow to continue to encourage explicit API design
- Extensions should be easy to explain, document, and understand
- Extensions should ideally still be a subset of runtime behaviour (if it works in a stub, it works
at runtime)
- Extensions should ideally not make `.pyi` files more permissive than `.py` files (if it works in a
stub, it works in an impl)
- **dot re-exports**: `from . import a` in an `__init__.pyi` is considered a re-export of `a`
(equivalent to `from . import a as a`). This is required to properly handle many stubs in the
wild. Currently it must be *exactly* `from . import ...`.
To that end we define the following extension:
> If an `__init__.pyi` for `mypackage` contains a `from...import` targetting a direct submodule of
> `mypackage`, then that submodule should be available as an attribute of `mypackage`.
Note: almost all tests in here have a stub and non-stub version, because we're interested in both
defining symbols *at all* and re-exporting them.
## Relative `from` Import of Direct Submodule in `__init__`
The `from . import submodule` idiom in an `__init__.pyi` is fairly explicit and we should definitely
support it.
We consider the `from . import submodule` idiom in an `__init__.pyi` an explicit re-export.
### In Stub
`mypackage/__init__.pyi`:
@ -63,7 +63,7 @@ reveal_type(mypackage.imported.X) # revealed: int
reveal_type(mypackage.fails.Y) # revealed: Unknown
```
## Relative `from` Import of Direct Submodule in `__init__` (Non-Stub Check)
### In Non-Stub
`mypackage/__init__.py`:
@ -95,8 +95,11 @@ reveal_type(mypackage.fails.Y) # revealed: Unknown
## Absolute `from` Import of Direct Submodule in `__init__`
If an absolute `from...import` happens to import a submodule, it works just as well as a relative
one.
If an absolute `from...import` happens to import a submodule (i.e. it's equivalent to
`from . import y`) we do not treat it as a re-export. We could, but we don't. (This is an arbitrary
decision and can be changed!)
### In Stub
`mypackage/__init__.pyi`:
@ -121,12 +124,14 @@ Y: int = 47
```py
import mypackage
reveal_type(mypackage.imported.X) # revealed: int
# TODO: this could work and would be nice to have?
# error: "has no member `imported`"
reveal_type(mypackage.imported.X) # revealed: Unknown
# error: "has no member `fails`"
reveal_type(mypackage.fails.Y) # revealed: Unknown
```
## Absolute `from` Import of Direct Submodule in `__init__` (Non-Stub Check)
### In Non-Stub
`mypackage/__init__.py`:
@ -159,7 +164,9 @@ reveal_type(mypackage.fails.Y) # revealed: Unknown
## Import of Direct Submodule in `__init__`
An `import` that happens to import a submodule does not expose the submodule as an attribute. (This
is an arbitrary decision and can be changed easily!)
is an arbitrary decision and can be changed!)
### In Stub
`mypackage/__init__.pyi`:
@ -178,12 +185,12 @@ X: int = 42
```py
import mypackage
# TODO: this is probably safe to allow, as it's an unambiguous import of a submodule
# TODO: this could work and would be nice to have?
# error: "has no member `imported`"
reveal_type(mypackage.imported.X) # revealed: Unknown
```
## Import of Direct Submodule in `__init__` (Non-Stub Check)
### In Non-Stub
`mypackage/__init__.py`:
@ -202,15 +209,17 @@ X: int = 42
```py
import mypackage
# TODO: this is probably safe to allow, as it's an unambiguous import of a submodule
# TODO: this could work and would be nice to have
# error: "has no member `imported`"
reveal_type(mypackage.imported.X) # revealed: Unknown
```
## Relative `from` Import of Nested Submodule in `__init__`
`from .submodule import nested` in an `__init__.pyi` is currently not supported as a way to expose
`mypackage.submodule` or `mypackage.submodule.nested` but it could be.
`from .submodule import nested` in an `__init__.pyi` does not re-export `mypackage.submodule`,
`mypackage.submodule.nested`, or `nested`.
### In Stub
`mypackage/__init__.pyi`:
@ -234,16 +243,21 @@ X: int = 42
```py
import mypackage
# TODO: this would be nice to allow
# error: "has no member `submodule`"
reveal_type(mypackage.submodule) # revealed: Unknown
# error: "has no member `submodule`"
reveal_type(mypackage.submodule.nested) # revealed: Unknown
# error: "has no member `submodule`"
reveal_type(mypackage.submodule.nested.X) # revealed: Unknown
# error: "has no member `nested`"
reveal_type(mypackage.nested) # revealed: Unknown
# error: "has no member `nested`"
reveal_type(mypackage.nested.X) # revealed: Unknown
```
## Relative `from` Import of Nested Submodule in `__init__` (Non-Stub Check)
### In Non-Stub
`from .submodule import nested` in an `__init__.py` exposes `mypackage.submodule` and `nested`.
`mypackage/__init__.py`:
@ -267,19 +281,22 @@ X: int = 42
```py
import mypackage
reveal_type(mypackage.submodule) # revealed: <module 'mypackage.submodule'>
# TODO: this would be nice to support
# error: "has no member `submodule`"
reveal_type(mypackage.submodule) # revealed: Unknown
# error: "has no member `submodule`"
# error: "has no member `nested`"
reveal_type(mypackage.submodule.nested) # revealed: Unknown
# error: "has no member `submodule`"
# error: "has no member `nested`"
reveal_type(mypackage.submodule.nested.X) # revealed: Unknown
reveal_type(mypackage.nested) # revealed: <module 'mypackage.submodule.nested'>
reveal_type(mypackage.nested.X) # revealed: int
```
## Absolute `from` Import of Nested Submodule in `__init__`
`from mypackage.submodule import nested` in an `__init__.pyi` is currently not supported as a way to
expose `mypackage.submodule` or `mypackage.submodule.nested` but it could be.
`from mypackage.submodule import nested` in an `__init__.pyi` does not re-export
`mypackage.submodule`, `mypackage.submodule.nested`, or `nested`.
### In Stub
`mypackage/__init__.pyi`:
@ -303,16 +320,22 @@ X: int = 42
```py
import mypackage
# TODO: this would be nice to support
# TODO: this could work and would be nice to have
# error: "has no member `submodule`"
reveal_type(mypackage.submodule) # revealed: Unknown
# error: "has no member `submodule`"
reveal_type(mypackage.submodule.nested) # revealed: Unknown
# error: "has no member `submodule`"
reveal_type(mypackage.submodule.nested.X) # revealed: Unknown
# error: "has no member `nested`"
reveal_type(mypackage.nested) # revealed: Unknown
# error: "has no member `nested`"
reveal_type(mypackage.nested.X) # revealed: Unknown
```
## Absolute `from` Import of Nested Submodule in `__init__` (Non-Stub Check)
### In Non-Stub
`from mypackage.submodule import nested` in an `__init__.py` only creates `nested`.
`mypackage/__init__.py`:
@ -343,12 +366,16 @@ reveal_type(mypackage.submodule) # revealed: Unknown
reveal_type(mypackage.submodule.nested) # revealed: Unknown
# error: "has no member `submodule`"
reveal_type(mypackage.submodule.nested.X) # revealed: Unknown
reveal_type(mypackage.nested) # revealed: <module 'mypackage.submodule.nested'>
reveal_type(mypackage.nested.X) # revealed: int
```
## Import of Nested Submodule in `__init__`
`import mypackage.submodule.nested` in an `__init__.pyi` is currently not supported as a way to
expose `mypackage.submodule` or `mypackage.submodule.nested` but it could be.
`import mypackage.submodule.nested` in an `__init__.pyi` does not re-export `mypackage.submodule` or
`mypackage.submodule.nested`.
### In Stub
`mypackage/__init__.pyi`:
@ -372,7 +399,6 @@ X: int = 42
```py
import mypackage
# TODO: this would be nice to support, and is probably safe to do as it's unambiguous
# error: "has no member `submodule`"
reveal_type(mypackage.submodule) # revealed: Unknown
# error: "has no member `submodule`"
@ -381,7 +407,10 @@ reveal_type(mypackage.submodule.nested) # revealed: Unknown
reveal_type(mypackage.submodule.nested.X) # revealed: Unknown
```
## Import of Nested Submodule in `__init__` (Non-Stub Check)
### In Non-Stub
`import mypackage.submodule.nested` in an `__init__.py` does not define `mypackage.submodule` or
`mypackage.submodule.nested` outside the package.
`mypackage/__init__.py`:
@ -405,7 +434,7 @@ X: int = 42
```py
import mypackage
# TODO: this would be nice to support, and is probably safe to do as it's unambiguous
# TODO: this would be nice to support
# error: "has no member `submodule`"
reveal_type(mypackage.submodule) # revealed: Unknown
# error: "has no member `submodule`"
@ -418,6 +447,8 @@ reveal_type(mypackage.submodule.nested.X) # revealed: Unknown
Renaming the submodule to something else disables the `__init__.pyi` idiom.
### In Stub
`mypackage/__init__.pyi`:
```pyi
@ -441,7 +472,7 @@ reveal_type(mypackage.imported.X) # revealed: Unknown
reveal_type(mypackage.imported_m.X) # revealed: Unknown
```
## Relative `from` Import of Direct Submodule in `__init__`, Mismatched Alias (Non-Stub Check)
### In Non-Stub
`mypackage/__init__.py`:
@ -471,6 +502,8 @@ reveal_type(mypackage.imported_m.X) # revealed: int
The `__init__.pyi` idiom should definitely always work if the submodule is renamed to itself, as
this is the re-export idiom.
### In Stub
`mypackage/__init__.pyi`:
```pyi
@ -491,7 +524,7 @@ import mypackage
reveal_type(mypackage.imported.X) # revealed: int
```
## Relative `from` Import of Direct Submodule in `__init__`, Matched Alias (Non-Stub Check)
### In Non-Stub
`mypackage/__init__.py`:
@ -518,6 +551,8 @@ reveal_type(mypackage.imported.X) # revealed: int
Even if the `__init__` idiom is in effect, star imports do not pick it up. (This is an arbitrary
decision that mostly fell out of the implementation details and can be changed!)
### In Stub
`mypackage/__init__.pyi`:
```pyi
@ -536,13 +571,13 @@ X: int = 42
```py
from mypackage import *
# TODO: this would be nice to support (available_submodule_attributes isn't visible to `*` imports)
# TODO: this would be nice to support
# error: "`imported` used when not defined"
reveal_type(imported.X) # revealed: Unknown
reveal_type(Z) # revealed: int
```
## Star Import Unaffected (Non-Stub Check)
### In Non-Stub
`mypackage/__init__.py`:
@ -569,9 +604,10 @@ reveal_type(Z) # revealed: int
## `from` Import of Non-Submodule
A from import that terminates in a non-submodule should not expose the intermediate submodules as
attributes. This is an arbitrary decision but on balance probably safe and correct, as otherwise it
would be hard for a stub author to be intentional about the submodules being exposed as attributes.
A `from` import that imports a non-submodule isn't currently a special case here (various
proposed/tested approaches did treat this specially).
### In Stub
`mypackage/__init__.pyi`:
@ -590,11 +626,11 @@ X: int = 42
```py
import mypackage
# error: "has no member `imported`"
# error: "no member `imported`"
reveal_type(mypackage.imported.X) # revealed: Unknown
```
## `from` Import of Non-Submodule (Non-Stub Check)
### In Non-Stub
`mypackage/__init__.py`:
@ -613,9 +649,7 @@ X: int = 42
```py
import mypackage
# TODO: this would be nice to support, as it works at runtime
# error: "has no member `imported`"
reveal_type(mypackage.imported.X) # revealed: Unknown
reveal_type(mypackage.imported.X) # revealed: int
```
## `from` Import of Other Package's Submodule
@ -623,6 +657,8 @@ reveal_type(mypackage.imported.X) # revealed: Unknown
`from mypackage import submodule` from outside the package is not modeled as a side-effect on
`mypackage`, even in the importing file (this could be changed!).
### In Stub
`mypackage/__init__.pyi`:
```pyi
@ -641,12 +677,13 @@ import mypackage
from mypackage import imported
# TODO: this would be nice to support, but it's dangerous with available_submodule_attributes
# for details, see: https://github.com/astral-sh/ty/issues/1488
reveal_type(imported.X) # revealed: int
# error: "has no member `imported`"
reveal_type(mypackage.imported.X) # revealed: Unknown
```
## `from` Import of Other Package's Submodule (Non-Stub Check)
### In Non-Stub
`mypackage/__init__.py`:
@ -676,6 +713,8 @@ reveal_type(mypackage.imported.X) # revealed: Unknown
`from . import submodule` from a sibling module is not modeled as a side-effect on `mypackage` or a
re-export from `submodule`.
### In Stub
`mypackage/__init__.pyi`:
```pyi
@ -707,7 +746,7 @@ reveal_type(imported.fails.Y) # revealed: Unknown
reveal_type(mypackage.fails.Y) # revealed: Unknown
```
## `from` Import of Sibling Module (Non-Stub Check)
### In Non-Stub
`mypackage/__init__.py`:
@ -752,9 +791,11 @@ Can easily result in the typechecker getting "confused" and thinking imports of
top-level package are referring to the subpackage and not the function/class. This issue can be
found with the `lobpcg` function in `scipy.sparse.linalg`.
This kind of failure mode is why the rule is restricted to *direct* submodule imports, as anything
more powerful than that in the current implementation strategy quickly gets the functions and
submodules mixed up.
We avoid this by ensuring that the imported name (the right-hand `funcmod` in
`from .funcmod import funcmod`) overwrites the submodule attribute (the left-hand `funcmod`), as it
does at runtime.
### In Stub
`mypackage/__init__.pyi`:
@ -788,7 +829,7 @@ from mypackage import funcmod
x = funcmod(1)
```
## Fractal Re-export Nameclash Problems (Non-Stub Check)
### In Non-Stub
`mypackage/__init__.py`:
@ -822,3 +863,311 @@ from mypackage import funcmod
x = funcmod(1)
```
## Re-export Nameclash Problems In Functions
`from` imports in an `__init__.py` at file scope should be visible to functions defined in the file:
`mypackage/__init__.py`:
```py
from .funcmod import funcmod
funcmod(1)
def run():
funcmod(2)
```
`mypackage/funcmod.py`:
```py
def funcmod(x: int) -> int:
return x
```
## Re-export Nameclash Problems In Try-Blocks
`from` imports in an `__init__.py` at file scope in a `try` block should be visible to functions
defined in the `try` block (regression test for a bug):
`mypackage/__init__.py`:
```py
try:
from .funcmod import funcmod
funcmod(1)
def run():
# TODO: this is a bug in how we analyze try-blocks
# error: [call-non-callable]
funcmod(2)
finally:
x = 1
```
`mypackage/funcmod.py`:
```py
def funcmod(x: int) -> int:
return x
```
## RHS `from` Imports In Functions
If a `from` import occurs in a function, the RHS symbols should only be visible in that function.
`mypackage/__init__.py`:
```py
def run1():
from .funcmod import funcmod
funcmod(1)
def run2():
from .funcmod import funcmod
funcmod(2)
def run3():
# error: [unresolved-reference]
funcmod(3)
# error: [unresolved-reference]
funcmod(4)
```
`mypackage/funcmod.py`:
```py
def funcmod(x: int) -> int:
return x
```
## LHS `from` Imports In Functions
If a `from` import occurs in a function, LHS symbols should only be visible in that function. This
very blatantly is not runtime-accurate, but exists to try to force you to write "obviously
deterministically correct" imports instead of relying on execution order.
`mypackage/__init__.py`:
```py
def run1():
from .funcmod import other
funcmod.funcmod(1)
def run2():
from .funcmod import other
# TODO: this is just a bug! We only register the first
# import of `funcmod` in the entire file, and not per-scope!
# error: [unresolved-reference]
funcmod.funcmod(2)
def run3():
# error: [unresolved-reference]
funcmod.funcmod(3)
# error: [unresolved-reference]
funcmod.funcmod(4)
```
`mypackage/funcmod.py`:
```py
other: int = 1
def funcmod(x: int) -> int:
return x
```
## LHS `from` Imports Overwrite Locals
The LHS of a `from..import` introduces a local symbol that overwrites any local with the same name.
This reflects actual runtime behaviour, although we're kinda assuming it hasn't been imported
already.
`mypackage/__init__.py`:
```py
funcmod = 0
from .funcmod import funcmod
funcmod(1)
```
`mypackage/funcmod.py`:
```py
def funcmod(x: int) -> int:
return x
```
## LHS `from` Imports Overwritten By Local Function
The LHS of a `from..import` introduces a local symbol that can be overwritten by defining a function
(or class) with the same name.
### In Stub
`mypackage/__init__.pyi`:
```pyi
from .funcmod import other
def funcmod(x: int) -> int: ...
```
`mypackage/funcmod/__init__.pyi`:
```pyi
def other(int) -> int: ...
```
`main.py`:
```py
from mypackage import funcmod
x = funcmod(1)
```
### In Non-Stub
`mypackage/__init__.py`:
```py
from .funcmod import other
def funcmod(x: int) -> int:
return x
```
`mypackage/funcmod/__init__.py`:
```py
def other(x: int) -> int:
return x
```
`main.py`:
```py
from mypackage import funcmod
x = funcmod(1)
```
## LHS `from` Imports Overwritten By Local Assignment
The LHS of a `from..import` introduces a local symbol that can be overwritten by assigning to it.
### In Stub
`mypackage/__init__.pyi`:
```pyi
from .funcmod import other
funcmod = other
```
`mypackage/funcmod/__init__.pyi`:
```pyi
def other(x: int) -> int: ...
```
`main.py`:
```py
from mypackage import funcmod
x = funcmod(1)
```
### In Non-Stub
`mypackage/__init__.py`:
```py
from .funcmod import other
funcmod = other
```
`mypackage/funcmod/__init__.py`:
```py
def other(x: int) -> int:
return x
```
`main.py`:
```py
from mypackage import funcmod
x = funcmod(1)
```
## LHS `from` Imports Only Apply The First Time
The LHS of a `from..import` of a submodule introduces a local symbol only the first time it
introduces a direct submodule. The second time does nothing.
### In Stub
`mypackage/__init__.pyi`:
```pyi
from .funcmod import funcmod as funcmod
from .funcmod import other
```
`mypackage/funcmod/__init__.pyi`:
```pyi
def other(x: int) -> int: ...
def funcmod(x: int) -> int: ...
```
`main.py`:
```py
from mypackage import funcmod
x = funcmod(1)
```
### In Non-Stub
`mypackage/__init__.py`:
```py
from .funcmod import funcmod
from .funcmod import other
```
`mypackage/funcmod/__init__.py`:
```py
def other(x: int) -> int:
return x
def funcmod(x: int) -> int:
return x
```
`main.py`:
```py
from mypackage import funcmod
x = funcmod(1)
```

View file

@ -607,23 +607,33 @@ module:
`module2.py`:
```py
import importlib
import importlib.abc
import imported
import imported.abc
```
`imported/__init__.pyi`:
```pyi
```
`imported/abc.pyi`:
```pyi
```
`main2.py`:
```py
import importlib
from module2 import importlib as other_importlib
import imported
from module2 import imported as other_imported
from ty_extensions import TypeOf, static_assert, is_equivalent_to
# error: [unresolved-attribute] "Module `importlib` has no member `abc`"
reveal_type(importlib.abc) # revealed: Unknown
# error: [unresolved-attribute] "Module `imported` has no member `abc`"
reveal_type(imported.abc) # revealed: Unknown
reveal_type(other_importlib.abc) # revealed: <module 'importlib.abc'>
reveal_type(other_imported.abc) # revealed: <module 'imported.abc'>
static_assert(not is_equivalent_to(TypeOf[importlib], TypeOf[other_importlib]))
static_assert(not is_equivalent_to(TypeOf[imported], TypeOf[other_imported]))
```
[materializations]: https://typing.python.org/en/latest/spec/glossary.html#term-materialize

View file

@ -6,12 +6,12 @@ use ruff_db::parsed::parsed_module;
use ruff_index::{IndexSlice, IndexVec};
use ruff_python_ast::NodeIndex;
use ruff_python_ast::name::Name;
use ruff_python_parser::semantic_errors::SemanticSyntaxError;
use rustc_hash::{FxHashMap, FxHashSet};
use salsa::Update;
use salsa::plumbing::AsId;
use crate::Db;
use crate::module_name::ModuleName;
use crate::node_key::NodeKey;
use crate::semantic_index::ast_ids::AstIds;
@ -28,7 +28,6 @@ use crate::semantic_index::scope::{
use crate::semantic_index::symbol::ScopedSymbolId;
use crate::semantic_index::use_def::{EnclosingSnapshotKey, ScopedEnclosingSnapshotId, UseDefMap};
use crate::semantic_model::HasTrackedScope;
use crate::{Db, Module, resolve_module};
pub mod ast_ids;
mod builder;
@ -84,65 +83,6 @@ pub(crate) fn imported_modules<'db>(db: &'db dyn Db, file: File) -> Arc<FxHashSe
semantic_index(db, file).imported_modules.clone()
}
/// Returns the set of relative submodules that are explicitly imported anywhere in
/// `importing_module`.
///
/// This set only considers `from...import` statements (but it could also include `import`).
/// It also only returns a non-empty result for `__init__.pyi` files.
/// See [`ModuleLiteralType::available_submodule_attributes`] for discussion
/// of why this analysis is intentionally limited.
///
/// This function specifically implements the rule that if an `__init__.pyi` file
/// contains a `from...import` that imports a direct submodule of the package,
/// that submodule should be available as an attribute of the package.
///
/// While we endeavour to accurately model import side-effects for `.py` files, we intentionally
/// limit them for `.pyi` files to encourage more intentional API design. The standard escape
/// hatches for this are the `import x as x` idiom or listing them in `__all__`, but in practice
/// some other idioms are popular.
///
/// In particular, many packages have their `__init__` include lines like
/// `from . import subpackage`, with the intent that `mypackage.subpackage` should be
/// available for anyone who only does `import mypackage`.
#[salsa::tracked(returns(deref), heap_size=ruff_memory_usage::heap_size)]
pub(crate) fn imported_relative_submodules_of_stub_package<'db>(
db: &'db dyn Db,
importing_module: Module<'db>,
) -> Box<[ModuleName]> {
let Some(file) = importing_module.file(db) else {
return Box::default();
};
if !file.is_package_stub(db) {
return Box::default();
}
semantic_index(db, file)
.maybe_imported_modules
.iter()
.filter_map(|import| {
let mut submodule = ModuleName::from_identifier_parts(
db,
file,
import.from_module.as_deref(),
import.level,
)
.ok()?;
// We only actually care if this is a direct submodule of the package
// so this part should actually be exactly the importing module.
let importing_module_name = importing_module.name(db);
if importing_module_name != &submodule {
return None;
}
submodule.extend(&ModuleName::new(import.submodule.as_str())?);
// Throw out the result if this doesn't resolve to an actual module.
// This is quite expensive, but we've gone through a lot of hoops to
// get here so it won't happen too much.
resolve_module(db, &submodule)?;
// Return only the relative part
submodule.relative_to(importing_module_name)
})
.collect()
}
/// Returns the use-def map for a specific `scope`.
///
/// Using [`use_def_map`] over [`semantic_index`] has the advantage that
@ -284,9 +224,6 @@ pub(crate) struct SemanticIndex<'db> {
/// The set of modules that are imported anywhere within this file.
imported_modules: Arc<FxHashSet<ModuleName>>,
/// `from...import` statements within this file that might import a submodule.
maybe_imported_modules: FxHashSet<MaybeModuleImport>,
/// Flags about the global scope (code usage impacting inference)
has_future_annotations: bool,
@ -300,16 +237,6 @@ pub(crate) struct SemanticIndex<'db> {
generator_functions: FxHashSet<FileScopeId>,
}
/// A `from...import` that may be an import of a module
///
/// Later analysis will determine if it is.
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, get_size2::GetSize)]
pub(crate) struct MaybeModuleImport {
level: u32,
from_module: Option<Name>,
submodule: Name,
}
impl<'db> SemanticIndex<'db> {
/// Returns the place table for a specific scope.
///

View file

@ -26,8 +26,8 @@ use crate::semantic_index::definition::{
AnnotatedAssignmentDefinitionNodeRef, AssignmentDefinitionNodeRef,
ComprehensionDefinitionNodeRef, Definition, DefinitionCategory, DefinitionNodeKey,
DefinitionNodeRef, Definitions, ExceptHandlerDefinitionNodeRef, ForStmtDefinitionNodeRef,
ImportDefinitionNodeRef, ImportFromDefinitionNodeRef, MatchPatternDefinitionNodeRef,
StarImportDefinitionNodeRef, WithItemDefinitionNodeRef,
ImportDefinitionNodeRef, ImportFromDefinitionNodeRef, ImportFromSubmoduleDefinitionNodeRef,
MatchPatternDefinitionNodeRef, StarImportDefinitionNodeRef, WithItemDefinitionNodeRef,
};
use crate::semantic_index::expression::{Expression, ExpressionKind};
use crate::semantic_index::place::{PlaceExpr, PlaceTableBuilder, ScopedPlaceId};
@ -47,9 +47,7 @@ use crate::semantic_index::symbol::{ScopedSymbolId, Symbol};
use crate::semantic_index::use_def::{
EnclosingSnapshotKey, FlowSnapshot, ScopedEnclosingSnapshotId, UseDefMapBuilder,
};
use crate::semantic_index::{
ExpressionsScopeMap, MaybeModuleImport, SemanticIndex, VisibleAncestorsIter,
};
use crate::semantic_index::{ExpressionsScopeMap, SemanticIndex, VisibleAncestorsIter};
use crate::semantic_model::HasTrackedScope;
use crate::unpack::{EvaluationMode, Unpack, UnpackKind, UnpackPosition, UnpackValue};
use crate::{Db, Program};
@ -113,7 +111,7 @@ pub(super) struct SemanticIndexBuilder<'db, 'ast> {
definitions_by_node: FxHashMap<DefinitionNodeKey, Definitions<'db>>,
expressions_by_node: FxHashMap<ExpressionNodeKey, Expression<'db>>,
imported_modules: FxHashSet<ModuleName>,
maybe_imported_modules: FxHashSet<MaybeModuleImport>,
seen_submodule_imports: FxHashSet<String>,
/// Hashset of all [`FileScopeId`]s that correspond to [generator functions].
///
/// [generator functions]: https://docs.python.org/3/glossary.html#term-generator
@ -151,7 +149,7 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
definitions_by_node: FxHashMap::default(),
expressions_by_node: FxHashMap::default(),
maybe_imported_modules: FxHashSet::default(),
seen_submodule_imports: FxHashSet::default(),
imported_modules: FxHashSet::default(),
generator_functions: FxHashSet::default(),
@ -1266,7 +1264,6 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
self.scopes_by_node.shrink_to_fit();
self.generator_functions.shrink_to_fit();
self.enclosing_snapshots.shrink_to_fit();
self.maybe_imported_modules.shrink_to_fit();
SemanticIndex {
place_tables,
@ -1279,7 +1276,6 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
scopes_by_node: self.scopes_by_node,
use_def_maps,
imported_modules: Arc::new(self.imported_modules),
maybe_imported_modules: self.maybe_imported_modules,
has_future_annotations: self.has_future_annotations,
enclosing_snapshots: self.enclosing_snapshots,
semantic_syntax_errors: self.semantic_syntax_errors.into_inner(),
@ -1453,6 +1449,43 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
self.current_use_def_map_mut()
.record_node_reachability(NodeKey::from_node(node));
// If we see:
//
// * `from .x.y import z` (must be relative!)
// * And we are in an `__init__.py(i)` (hereafter `thispackage`)
// * And this is the first time we've seen `from .x` in this module
//
// We introduce a local definition `x = <module 'thispackage.x'>` that occurs
// before the `z = ...` declaration the import introduces. This models the fact
// that the *first* time that you import 'thispackage.x' the python runtime creates
// `x` as a variable in the global scope of `thispackage`.
//
// This is not a perfect simulation of actual runtime behaviour for *various*
// reasons but it works well for most practical purposes. In particular it's nice
// that `x` can be freely overwritten, and that we don't assume that an import
// in one function is visible in another function.
//
// TODO: Also support `from thispackage.x.y import z`?
// TODO: `seen_submodule_imports` should be per-scope and not per-file
// (if two functions import `.x`, they both should believe `x` is defined)
if node.level == 1
&& let Some(submodule) = &node.module
&& let Some(parsed_submodule) = ModuleName::new(submodule.as_str())
&& let Some(direct_submodule) = parsed_submodule.components().next()
&& self.file.is_package(self.db)
&& !self.seen_submodule_imports.contains(direct_submodule)
{
self.seen_submodule_imports
.insert(direct_submodule.to_owned());
let direct_submodule_name = Name::new(direct_submodule);
let symbol = self.add_symbol(direct_submodule_name);
self.add_definition(
symbol.into(),
ImportFromSubmoduleDefinitionNodeRef { node, submodule },
);
}
let mut found_star = false;
for (alias_index, alias) in node.names.iter().enumerate() {
if &alias.name == "*" {
@ -1559,20 +1592,18 @@ impl<'ast> Visitor<'ast> for SemanticIndexBuilder<'_, 'ast> {
}
let (symbol_name, is_reexported) = if let Some(asname) = &alias.asname {
// It's re-exported if it's `from ... import x as x`
(&asname.id, asname.id == alias.name.id)
} else {
(&alias.name.id, false)
// It's re-exported if it's `from . import x` in an `__init__.pyi`
(
&alias.name.id,
node.level == 1
&& node.module.is_none()
&& self.file.is_package(self.db),
)
};
// If there's no alias or a redundant alias, record this as a potential import of a submodule
if alias.asname.is_none() || is_reexported {
self.maybe_imported_modules.insert(MaybeModuleImport {
level: node.level,
from_module: node.module.clone().map(Into::into),
submodule: alias.name.clone().into(),
});
}
// Look for imports `from __future__ import annotations`, ignore `as ...`
// We intentionally don't enforce the rules about location of `__future__`
// imports here, we assume the user's intent was to apply the `__future__`

View file

@ -3,6 +3,7 @@ use std::ops::Deref;
use ruff_db::files::{File, FileRange};
use ruff_db::parsed::{ParsedModuleRef, parsed_module};
use ruff_python_ast as ast;
use ruff_python_ast::name::Name;
use ruff_text_size::{Ranged, TextRange};
use crate::Db;
@ -209,6 +210,7 @@ impl<'db> DefinitionState<'db> {
pub(crate) enum DefinitionNodeRef<'ast, 'db> {
Import(ImportDefinitionNodeRef<'ast>),
ImportFrom(ImportFromDefinitionNodeRef<'ast>),
ImportFromSubmodule(ImportFromSubmoduleDefinitionNodeRef<'ast>),
ImportStar(StarImportDefinitionNodeRef<'ast>),
For(ForStmtDefinitionNodeRef<'ast, 'db>),
Function(&'ast ast::StmtFunctionDef),
@ -290,6 +292,12 @@ impl<'ast> From<ImportFromDefinitionNodeRef<'ast>> for DefinitionNodeRef<'ast, '
}
}
impl<'ast> From<ImportFromSubmoduleDefinitionNodeRef<'ast>> for DefinitionNodeRef<'ast, '_> {
fn from(node_ref: ImportFromSubmoduleDefinitionNodeRef<'ast>) -> Self {
Self::ImportFromSubmodule(node_ref)
}
}
impl<'ast, 'db> From<ForStmtDefinitionNodeRef<'ast, 'db>> for DefinitionNodeRef<'ast, 'db> {
fn from(value: ForStmtDefinitionNodeRef<'ast, 'db>) -> Self {
Self::For(value)
@ -357,7 +365,11 @@ pub(crate) struct ImportFromDefinitionNodeRef<'ast> {
pub(crate) alias_index: usize,
pub(crate) is_reexported: bool,
}
#[derive(Copy, Clone, Debug)]
pub(crate) struct ImportFromSubmoduleDefinitionNodeRef<'ast> {
pub(crate) node: &'ast ast::StmtImportFrom,
pub(crate) submodule: &'ast ast::Identifier,
}
#[derive(Copy, Clone, Debug)]
pub(crate) struct AssignmentDefinitionNodeRef<'ast, 'db> {
pub(crate) unpack: Option<(UnpackPosition, Unpack<'db>)>,
@ -427,7 +439,6 @@ impl<'db> DefinitionNodeRef<'_, 'db> {
alias_index,
is_reexported,
}),
DefinitionNodeRef::ImportFrom(ImportFromDefinitionNodeRef {
node,
alias_index,
@ -437,6 +448,13 @@ impl<'db> DefinitionNodeRef<'_, 'db> {
alias_index,
is_reexported,
}),
DefinitionNodeRef::ImportFromSubmodule(ImportFromSubmoduleDefinitionNodeRef {
node,
submodule,
}) => DefinitionKind::ImportFromSubmodule(ImportFromSubmoduleDefinitionKind {
node: AstNodeRef::new(parsed, node),
submodule: submodule.as_str().into(),
}),
DefinitionNodeRef::ImportStar(star_import) => {
let StarImportDefinitionNodeRef { node, symbol_id } = star_import;
DefinitionKind::StarImport(StarImportDefinitionKind {
@ -562,7 +580,10 @@ impl<'db> DefinitionNodeRef<'_, 'db> {
alias_index,
is_reexported: _,
}) => (&node.names[alias_index]).into(),
Self::ImportFromSubmodule(ImportFromSubmoduleDefinitionNodeRef {
node,
submodule: _,
}) => node.into(),
// INVARIANT: for an invalid-syntax statement such as `from foo import *, bar, *`,
// we only create a `StarImportDefinitionKind` for the *first* `*` alias in the names list.
Self::ImportStar(StarImportDefinitionNodeRef { node, symbol_id: _ }) => node
@ -661,6 +682,7 @@ impl DefinitionCategory {
pub enum DefinitionKind<'db> {
Import(ImportDefinitionKind),
ImportFrom(ImportFromDefinitionKind),
ImportFromSubmodule(ImportFromSubmoduleDefinitionKind),
StarImport(StarImportDefinitionKind),
Function(AstNodeRef<ast::StmtFunctionDef>),
Class(AstNodeRef<ast::StmtClassDef>),
@ -687,6 +709,7 @@ impl DefinitionKind<'_> {
match self {
DefinitionKind::Import(import) => import.is_reexported(),
DefinitionKind::ImportFrom(import) => import.is_reexported(),
DefinitionKind::ImportFromSubmodule(_) => false,
_ => true,
}
}
@ -704,6 +727,7 @@ impl DefinitionKind<'_> {
DefinitionKind::Import(_)
| DefinitionKind::ImportFrom(_)
| DefinitionKind::StarImport(_)
| DefinitionKind::ImportFromSubmodule(_)
)
}
@ -719,6 +743,7 @@ impl DefinitionKind<'_> {
match self {
DefinitionKind::Import(import) => import.alias(module).range(),
DefinitionKind::ImportFrom(import) => import.alias(module).range(),
DefinitionKind::ImportFromSubmodule(import) => import.import(module).range(),
DefinitionKind::StarImport(import) => import.alias(module).range(),
DefinitionKind::Function(function) => function.node(module).name.range(),
DefinitionKind::Class(class) => class.node(module).name.range(),
@ -756,6 +781,7 @@ impl DefinitionKind<'_> {
match self {
DefinitionKind::Import(import) => import.alias(module).range(),
DefinitionKind::ImportFrom(import) => import.alias(module).range(),
DefinitionKind::ImportFromSubmodule(import) => import.import(module).range(),
DefinitionKind::StarImport(import) => import.import(module).range(),
DefinitionKind::Function(function) => function.node(module).range(),
DefinitionKind::Class(class) => class.node(module).range(),
@ -846,6 +872,7 @@ impl DefinitionKind<'_> {
| DefinitionKind::Comprehension(_)
| DefinitionKind::WithItem(_)
| DefinitionKind::MatchPattern(_)
| DefinitionKind::ImportFromSubmodule(_)
| DefinitionKind::ExceptHandler(_) => DefinitionCategory::Binding,
}
}
@ -991,6 +1018,21 @@ impl ImportFromDefinitionKind {
self.is_reexported
}
}
#[derive(Clone, Debug, get_size2::GetSize)]
pub struct ImportFromSubmoduleDefinitionKind {
node: AstNodeRef<ast::StmtImportFrom>,
submodule: Name,
}
impl ImportFromSubmoduleDefinitionKind {
pub fn import<'ast>(&self, module: &'ast ParsedModuleRef) -> &'ast ast::StmtImportFrom {
self.node.node(module)
}
pub(crate) fn submodule(&self) -> &Name {
&self.submodule
}
}
#[derive(Clone, Debug, get_size2::GetSize)]
pub struct AssignmentDefinitionKind<'db> {
@ -1121,6 +1163,12 @@ impl From<&ast::Alias> for DefinitionNodeKey {
}
}
impl From<&ast::StmtImportFrom> for DefinitionNodeKey {
fn from(node: &ast::StmtImportFrom) -> Self {
Self(NodeKey::from_node(node))
}
}
impl From<&ast::StmtFunctionDef> for DefinitionNodeKey {
fn from(node: &ast::StmtFunctionDef) -> Self {
Self(NodeKey::from_node(node))

View file

@ -39,9 +39,7 @@ use crate::place::{
use crate::semantic_index::definition::{Definition, DefinitionKind};
use crate::semantic_index::place::ScopedPlaceId;
use crate::semantic_index::scope::ScopeId;
use crate::semantic_index::{
imported_modules, imported_relative_submodules_of_stub_package, place_table, semantic_index,
};
use crate::semantic_index::{imported_modules, place_table, semantic_index};
use crate::suppression::check_suppressions;
use crate::types::bound_super::BoundSuperType;
use crate::types::call::{Binding, Bindings, CallArguments, CallableBinding};
@ -11302,29 +11300,23 @@ impl<'db> ModuleLiteralType<'db> {
///
/// # Rules
///
/// We have two rules for whether a submodule attribute is defined:
/// Because of the excessive power and danger of this method, we currently have only one rule:
///
/// * If the importing file include `import x.y` then `x.y` is defined in the importing file.
/// This is an easy rule to justify because `import` can only ever import a module, and so
/// * If the importing file includes `import x.y` then `x.y` is defined in the importing file.
/// This is an easy rule to justify because `import` can only ever import a module, and the
/// only reason to do it is to explicitly introduce those submodules and attributes, so it
/// *should* shadow any non-submodule of the same name.
///
/// * If the module is an `__init__.pyi` for `mypackage`, and it contains a `from...import`
/// that normalizes to `from mypackage import submodule`, then `mypackage.submodule` is
/// defined in all files. This supports the `from . import submodule` idiom. Critically,
/// we do *not* allow `from mypackage.nested import submodule` to affect `mypackage`.
/// The idea here is that `from mypackage import submodule` *from mypackage itself* can
/// only ever reasonably be an import of a submodule. It doesn't make any sense to import
/// a function or class from yourself! (You *can* do it but... why? Don't? Please?)
/// `from x.y import z` instances are currently ignored because the `x.y` part may not be a
/// side-effect the user actually cares about, and the `z` component may not be a submodule.
///
/// We instead prefer handling most other import effects as definitions in the scope of
/// the current file (i.e. [`crate::semantic_index::definition::ImportFromDefinitionNodeRef`]).
fn available_submodule_attributes(&self, db: &'db dyn Db) -> impl Iterator<Item = Name> {
self.importing_file(db)
.into_iter()
.flat_map(|file| imported_modules(db, file))
.filter_map(|submodule_name| submodule_name.relative_to(self.module(db).name(db)))
.chain(
imported_relative_submodules_of_stub_package(db, self.module(db))
.iter()
.cloned(),
)
.filter_map(|relative_submodule| relative_submodule.components().next().map(Name::from))
}

View file

@ -1285,7 +1285,7 @@ mod resolve_definition {
let file = definition.file(db);
let module = parsed_module(db, file).load(db);
let import_node = import_from_def.import(&module);
let alias = import_from_def.alias(&module);
let name = &import_from_def.alias(&module).name;
// For `ImportFrom`, we need to resolve the original imported symbol name
// (alias.name), not the local alias (symbol_name)
@ -1293,7 +1293,7 @@ mod resolve_definition {
db,
file,
import_node,
&alias.name,
name,
visited,
alias_resolution,
)
@ -1625,6 +1625,7 @@ mod resolve_definition {
DefinitionKind::TypeAlias(_)
| DefinitionKind::Import(_)
| DefinitionKind::ImportFrom(_)
| DefinitionKind::ImportFromSubmodule(_)
| DefinitionKind::StarImport(_)
| DefinitionKind::NamedExpression(_)
| DefinitionKind::Assignment(_)

View file

@ -4,6 +4,7 @@ use itertools::{Either, Itertools};
use ruff_db::diagnostic::{Annotation, DiagnosticId, Severity};
use ruff_db::files::File;
use ruff_db::parsed::ParsedModuleRef;
use ruff_python_ast::name::Name;
use ruff_python_ast::visitor::{Visitor, walk_expr};
use ruff_python_ast::{
self as ast, AnyNodeRef, ExprContext, HasNodeIndex, NodeIndex, PythonVersion,
@ -1214,6 +1215,13 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
definition,
);
}
DefinitionKind::ImportFromSubmodule(import_from) => {
self.infer_import_from_submodule_definition(
import_from.import(self.module()),
import_from.submodule(),
definition,
);
}
DefinitionKind::StarImport(import) => {
self.infer_import_from_definition(
import.import(self.module()),
@ -5893,6 +5901,99 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
}
}
/// Infer the implicit local definition `x = <module 'thispackage.x'>` that
/// `from .x.y import z` can introduce in an `__init__.py(i)`.
///
/// For the definition `z`, see [`TypeInferenceBuilder::infer_import_from_definition`].
fn infer_import_from_submodule_definition(
&mut self,
import_from: &ast::StmtImportFrom,
submodule: &Name,
definition: Definition<'db>,
) {
// Although the *actual* runtime semantic of this kind of statement is to
// introduce a variable in the global scope of this module, we want to
// encourage users to write code that doesn't have dependence on execution-order.
//
// By introducing it as a local variable in the scope the import occurs in,
// we effectively require the developer to either do the import at the start of
// the file where it belongs, or to repeat the import in every function that
// wants to use it, which "definitely" works.
//
// (It doesn't actually "definitely" work because only the first import of `thispackage.x`
// will ever set `x`, and any subsequent overwrites of it will permanently clobber it.
// Also, a local variable `x` in a function should always shadow the submodule because
// the submodule is defined at file-scope. However, both of these issues are much more
// narrow, so this approach seems to work well in practice!)
// Get this package's module by resolving `.`
let Ok(module_name) = ModuleName::from_identifier_parts(self.db(), self.file(), None, 1)
else {
self.add_binding(import_from.into(), definition, |_, _| Type::unknown());
return;
};
let Some(module) = resolve_module(self.db(), &module_name) else {
self.add_binding(import_from.into(), definition, |_, _| Type::unknown());
return;
};
// Now construct the submodule `.x`
assert!(
!submodule.is_empty(),
"ImportFromSubmoduleDefinitionKind constructed with empty module"
);
let name = submodule
.split_once('.')
.map(|(first, _)| first)
.unwrap_or(submodule.as_str());
let full_submodule_name = ModuleName::new(name).map(|final_part| {
let mut ret = module_name.clone();
ret.extend(&final_part);
ret
});
// And try to import it
if let Some(submodule_type) = full_submodule_name
.as_ref()
.and_then(|submodule_name| self.module_type_from_name(submodule_name))
{
// Success, introduce a binding!
//
// We explicitly don't introduce a *declaration* because it's actual ok
// (and fairly common) to overwrite this import with a function or class
// and we don't want it to be a type error to do so.
self.add_binding(import_from.into(), definition, |_, _| submodule_type);
return;
}
// That didn't work, try to produce diagnostics
self.add_binding(import_from.into(), definition, |_, _| Type::unknown());
if !self.is_reachable(import_from) {
return;
}
let Some(builder) = self
.context
.report_lint(&UNRESOLVED_IMPORT, AnyNodeRef::StmtImportFrom(import_from))
else {
return;
};
let diagnostic = builder.into_diagnostic(format_args!(
"Module `{module_name}` has no submodule `{name}`"
));
if let Some(full_submodule_name) = full_submodule_name {
hint_if_stdlib_submodule_exists_on_other_versions(
self.db(),
diagnostic,
&full_submodule_name,
module,
);
}
}
fn infer_return_statement(&mut self, ret: &ast::StmtReturn) {
let tcx = if ret.value.is_some() {
nearest_enclosing_function(self.db(), self.index, self.scope())