mirror of
https://github.com/astral-sh/ruff
synced 2026-01-21 21:40:51 -05:00
Refactor import-tracking to leverage existing AST bindings (#1856)
This PR refactors our import-tracking logic to leverage our existing logic for tracking bindings. It's both a significant simplification, a significant improvement (as we can now track reassignments), and closes out a bunch of subtle bugs. Though the AST tracks all bindings (e.g., when parsing `import os as foo`, we bind the name `foo` to a `BindingKind::Importation` that points to the `os` module), when I went to implement import tracking (e.g., to ensure that if the user references `List`, it's actually `typing.List`), I added a parallel system specifically for this use-case. That was a mistake, for a few reasons: 1. It didn't track reassignments, so if you had `from typing import List`, but `List` was later overridden, we'd still consider any reference to `List` to be `typing.List`. 2. It required a bunch of extra logic, include complex logic to try and optimize the lookups, since it's such a hot codepath. 3. There were a few bugs in the implementation that were just hard to correct under the existing abstractions (e.g., if you did `from typing import Optional as Foo`, then we'd treat any reference to `Foo` _or_ `Optional` as `typing.Optional` (even though, in that case, `Optional` was really unbound). The new implementation goes through our existing binding tracking: when we get a reference, we find the appropriate binding given the current scope stack, and normalize it back to its original target. Closes #1690. Closes #1790.
This commit is contained in:
Charlie Marsh
GitHub
@@ -1,8 +1,8 @@
|
||||
use once_cell::sync::Lazy;
|
||||
use rustc_hash::{FxHashMap, FxHashSet};
|
||||
use rustc_hash::FxHashSet;
|
||||
use rustpython_ast::{Expr, ExprKind};
|
||||
|
||||
use crate::ast::helpers::{collect_call_paths, dealias_call_path, match_call_path};
|
||||
use crate::checkers::ast::Checker;
|
||||
|
||||
// See: https://pypi.org/project/typing-extensions/
|
||||
pub static TYPING_EXTENSIONS: Lazy<FxHashSet<&'static str>> = Lazy::new(|| {
|
||||
@@ -62,225 +62,190 @@ pub static TYPING_EXTENSIONS: Lazy<FxHashSet<&'static str>> = Lazy::new(|| {
|
||||
});
|
||||
|
||||
// See: https://docs.python.org/3/library/typing.html
|
||||
static SUBSCRIPTS: Lazy<FxHashMap<&'static str, Vec<&'static str>>> = Lazy::new(|| {
|
||||
let mut subscripts: FxHashMap<&'static str, Vec<&'static str>> = FxHashMap::default();
|
||||
for (module, name) in [
|
||||
// builtins
|
||||
("", "dict"),
|
||||
("", "frozenset"),
|
||||
("", "list"),
|
||||
("", "set"),
|
||||
("", "tuple"),
|
||||
("", "type"),
|
||||
// `collections`
|
||||
("collections", "ChainMap"),
|
||||
("collections", "Counter"),
|
||||
("collections", "OrderedDict"),
|
||||
("collections", "defaultdict"),
|
||||
("collections", "deque"),
|
||||
// `collections.abc`
|
||||
("collections.abc", "AsyncGenerator"),
|
||||
("collections.abc", "AsyncIterable"),
|
||||
("collections.abc", "AsyncIterator"),
|
||||
("collections.abc", "Awaitable"),
|
||||
("collections.abc", "ByteString"),
|
||||
("collections.abc", "Callable"),
|
||||
("collections.abc", "Collection"),
|
||||
("collections.abc", "Container"),
|
||||
("collections.abc", "Coroutine"),
|
||||
("collections.abc", "Generator"),
|
||||
("collections.abc", "ItemsView"),
|
||||
("collections.abc", "Iterable"),
|
||||
("collections.abc", "Iterator"),
|
||||
("collections.abc", "KeysView"),
|
||||
("collections.abc", "Mapping"),
|
||||
("collections.abc", "MappingView"),
|
||||
("collections.abc", "MutableMapping"),
|
||||
("collections.abc", "MutableSequence"),
|
||||
("collections.abc", "MutableSet"),
|
||||
("collections.abc", "Reversible"),
|
||||
("collections.abc", "Sequence"),
|
||||
("collections.abc", "Set"),
|
||||
("collections.abc", "ValuesView"),
|
||||
// `contextlib`
|
||||
("contextlib", "AbstractAsyncContextManager"),
|
||||
("contextlib", "AbstractContextManager"),
|
||||
// `re`
|
||||
("re", "Match"),
|
||||
("re", "Pattern"),
|
||||
// `typing`
|
||||
("typing", "AbstractSet"),
|
||||
("typing", "AsyncContextManager"),
|
||||
("typing", "AsyncGenerator"),
|
||||
("typing", "AsyncIterator"),
|
||||
("typing", "Awaitable"),
|
||||
("typing", "BinaryIO"),
|
||||
("typing", "ByteString"),
|
||||
("typing", "Callable"),
|
||||
("typing", "ChainMap"),
|
||||
("typing", "ClassVar"),
|
||||
("typing", "Collection"),
|
||||
("typing", "Concatenate"),
|
||||
("typing", "Container"),
|
||||
("typing", "ContextManager"),
|
||||
("typing", "Coroutine"),
|
||||
("typing", "Counter"),
|
||||
("typing", "DefaultDict"),
|
||||
("typing", "Deque"),
|
||||
("typing", "Dict"),
|
||||
("typing", "Final"),
|
||||
("typing", "FrozenSet"),
|
||||
("typing", "Generator"),
|
||||
("typing", "Generic"),
|
||||
("typing", "IO"),
|
||||
("typing", "ItemsView"),
|
||||
("typing", "Iterable"),
|
||||
("typing", "Iterator"),
|
||||
("typing", "KeysView"),
|
||||
("typing", "List"),
|
||||
("typing", "Mapping"),
|
||||
("typing", "Match"),
|
||||
("typing", "MutableMapping"),
|
||||
("typing", "MutableSequence"),
|
||||
("typing", "MutableSet"),
|
||||
("typing", "Optional"),
|
||||
("typing", "OrderedDict"),
|
||||
("typing", "Pattern"),
|
||||
("typing", "Reversible"),
|
||||
("typing", "Sequence"),
|
||||
("typing", "Set"),
|
||||
("typing", "TextIO"),
|
||||
("typing", "Tuple"),
|
||||
("typing", "Type"),
|
||||
("typing", "TypeGuard"),
|
||||
("typing", "Union"),
|
||||
("typing", "Unpack"),
|
||||
("typing", "ValuesView"),
|
||||
// `typing.io`
|
||||
("typing.io", "BinaryIO"),
|
||||
("typing.io", "IO"),
|
||||
("typing.io", "TextIO"),
|
||||
// `typing.re`
|
||||
("typing.re", "Match"),
|
||||
("typing.re", "Pattern"),
|
||||
// `typing_extensions`
|
||||
("typing_extensions", "AsyncContextManager"),
|
||||
("typing_extensions", "AsyncGenerator"),
|
||||
("typing_extensions", "AsyncIterable"),
|
||||
("typing_extensions", "AsyncIterator"),
|
||||
("typing_extensions", "Awaitable"),
|
||||
("typing_extensions", "ChainMap"),
|
||||
("typing_extensions", "ClassVar"),
|
||||
("typing_extensions", "Concatenate"),
|
||||
("typing_extensions", "ContextManager"),
|
||||
("typing_extensions", "Coroutine"),
|
||||
("typing_extensions", "Counter"),
|
||||
("typing_extensions", "DefaultDict"),
|
||||
("typing_extensions", "Deque"),
|
||||
("typing_extensions", "Type"),
|
||||
// `weakref`
|
||||
("weakref", "WeakKeyDictionary"),
|
||||
("weakref", "WeakSet"),
|
||||
("weakref", "WeakValueDictionary"),
|
||||
] {
|
||||
subscripts.entry(name).or_default().push(module);
|
||||
}
|
||||
subscripts
|
||||
});
|
||||
const SUBSCRIPTS: &[&[&str]] = &[
|
||||
// builtins
|
||||
&["", "dict"],
|
||||
&["", "frozenset"],
|
||||
&["", "list"],
|
||||
&["", "set"],
|
||||
&["", "tuple"],
|
||||
&["", "type"],
|
||||
// `collections`
|
||||
&["collections", "ChainMap"],
|
||||
&["collections", "Counter"],
|
||||
&["collections", "OrderedDict"],
|
||||
&["collections", "defaultdict"],
|
||||
&["collections", "deque"],
|
||||
// `collections.abc`
|
||||
&["collections", "abc", "AsyncGenerator"],
|
||||
&["collections", "abc", "AsyncIterable"],
|
||||
&["collections", "abc", "AsyncIterator"],
|
||||
&["collections", "abc", "Awaitable"],
|
||||
&["collections", "abc", "ByteString"],
|
||||
&["collections", "abc", "Callable"],
|
||||
&["collections", "abc", "Collection"],
|
||||
&["collections", "abc", "Container"],
|
||||
&["collections", "abc", "Coroutine"],
|
||||
&["collections", "abc", "Generator"],
|
||||
&["collections", "abc", "ItemsView"],
|
||||
&["collections", "abc", "Iterable"],
|
||||
&["collections", "abc", "Iterator"],
|
||||
&["collections", "abc", "KeysView"],
|
||||
&["collections", "abc", "Mapping"],
|
||||
&["collections", "abc", "MappingView"],
|
||||
&["collections", "abc", "MutableMapping"],
|
||||
&["collections", "abc", "MutableSequence"],
|
||||
&["collections", "abc", "MutableSet"],
|
||||
&["collections", "abc", "Reversible"],
|
||||
&["collections", "abc", "Sequence"],
|
||||
&["collections", "abc", "Set"],
|
||||
&["collections", "abc", "ValuesView"],
|
||||
// `contextlib`
|
||||
&["contextlib", "AbstractAsyncContextManager"],
|
||||
&["contextlib", "AbstractContextManager"],
|
||||
// `re`
|
||||
&["re", "Match"],
|
||||
&["re", "Pattern"],
|
||||
// `typing`
|
||||
&["typing", "AbstractSet"],
|
||||
&["typing", "AsyncContextManager"],
|
||||
&["typing", "AsyncGenerator"],
|
||||
&["typing", "AsyncIterator"],
|
||||
&["typing", "Awaitable"],
|
||||
&["typing", "BinaryIO"],
|
||||
&["typing", "ByteString"],
|
||||
&["typing", "Callable"],
|
||||
&["typing", "ChainMap"],
|
||||
&["typing", "ClassVar"],
|
||||
&["typing", "Collection"],
|
||||
&["typing", "Concatenate"],
|
||||
&["typing", "Container"],
|
||||
&["typing", "ContextManager"],
|
||||
&["typing", "Coroutine"],
|
||||
&["typing", "Counter"],
|
||||
&["typing", "DefaultDict"],
|
||||
&["typing", "Deque"],
|
||||
&["typing", "Dict"],
|
||||
&["typing", "Final"],
|
||||
&["typing", "FrozenSet"],
|
||||
&["typing", "Generator"],
|
||||
&["typing", "Generic"],
|
||||
&["typing", "IO"],
|
||||
&["typing", "ItemsView"],
|
||||
&["typing", "Iterable"],
|
||||
&["typing", "Iterator"],
|
||||
&["typing", "KeysView"],
|
||||
&["typing", "List"],
|
||||
&["typing", "Mapping"],
|
||||
&["typing", "Match"],
|
||||
&["typing", "MutableMapping"],
|
||||
&["typing", "MutableSequence"],
|
||||
&["typing", "MutableSet"],
|
||||
&["typing", "Optional"],
|
||||
&["typing", "OrderedDict"],
|
||||
&["typing", "Pattern"],
|
||||
&["typing", "Reversible"],
|
||||
&["typing", "Sequence"],
|
||||
&["typing", "Set"],
|
||||
&["typing", "TextIO"],
|
||||
&["typing", "Tuple"],
|
||||
&["typing", "Type"],
|
||||
&["typing", "TypeGuard"],
|
||||
&["typing", "Union"],
|
||||
&["typing", "Unpack"],
|
||||
&["typing", "ValuesView"],
|
||||
// `typing.io`
|
||||
&["typing", "io", "BinaryIO"],
|
||||
&["typing", "io", "IO"],
|
||||
&["typing", "io", "TextIO"],
|
||||
// `typing.re`
|
||||
&["typing", "re", "Match"],
|
||||
&["typing", "re", "Pattern"],
|
||||
// `typing_extensions`
|
||||
&["typing_extensions", "AsyncContextManager"],
|
||||
&["typing_extensions", "AsyncGenerator"],
|
||||
&["typing_extensions", "AsyncIterable"],
|
||||
&["typing_extensions", "AsyncIterator"],
|
||||
&["typing_extensions", "Awaitable"],
|
||||
&["typing_extensions", "ChainMap"],
|
||||
&["typing_extensions", "ClassVar"],
|
||||
&["typing_extensions", "Concatenate"],
|
||||
&["typing_extensions", "ContextManager"],
|
||||
&["typing_extensions", "Coroutine"],
|
||||
&["typing_extensions", "Counter"],
|
||||
&["typing_extensions", "DefaultDict"],
|
||||
&["typing_extensions", "Deque"],
|
||||
&["typing_extensions", "Type"],
|
||||
// `weakref`
|
||||
&["weakref", "WeakKeyDictionary"],
|
||||
&["weakref", "WeakSet"],
|
||||
&["weakref", "WeakValueDictionary"],
|
||||
];
|
||||
|
||||
// See: https://docs.python.org/3/library/typing.html
|
||||
static PEP_593_SUBSCRIPTS: Lazy<FxHashMap<&'static str, Vec<&'static str>>> = Lazy::new(|| {
|
||||
let mut subscripts: FxHashMap<&'static str, Vec<&'static str>> = FxHashMap::default();
|
||||
for (module, name) in [
|
||||
// `typing`
|
||||
("typing", "Annotated"),
|
||||
// `typing_extensions`
|
||||
("typing_extensions", "Annotated"),
|
||||
] {
|
||||
subscripts.entry(name).or_default().push(module);
|
||||
}
|
||||
subscripts
|
||||
});
|
||||
const PEP_593_SUBSCRIPTS: &[&[&str]] = &[
|
||||
// `typing`
|
||||
&["typing", "Annotated"],
|
||||
// `typing_extensions`
|
||||
&["typing_extensions", "Annotated"],
|
||||
];
|
||||
|
||||
pub enum SubscriptKind {
|
||||
AnnotatedSubscript,
|
||||
PEP593AnnotatedSubscript,
|
||||
}
|
||||
|
||||
pub fn match_annotated_subscript<'a, F>(
|
||||
expr: &Expr,
|
||||
from_imports: &FxHashMap<&str, FxHashSet<&str>>,
|
||||
import_aliases: &FxHashMap<&str, &str>,
|
||||
typing_modules: impl Iterator<Item = &'a str>,
|
||||
is_builtin: F,
|
||||
) -> Option<SubscriptKind>
|
||||
where
|
||||
F: Fn(&str) -> bool,
|
||||
{
|
||||
pub fn match_annotated_subscript(checker: &Checker, expr: &Expr) -> Option<SubscriptKind> {
|
||||
if !matches!(
|
||||
expr.node,
|
||||
ExprKind::Name { .. } | ExprKind::Attribute { .. }
|
||||
) {
|
||||
return None;
|
||||
}
|
||||
let call_path = dealias_call_path(collect_call_paths(expr), import_aliases);
|
||||
if let Some(member) = call_path.last() {
|
||||
if let Some(modules) = SUBSCRIPTS.get(member) {
|
||||
for module in modules {
|
||||
if match_call_path(&call_path, module, member, from_imports)
|
||||
&& (!module.is_empty() || is_builtin(member))
|
||||
{
|
||||
return Some(SubscriptKind::AnnotatedSubscript);
|
||||
|
||||
checker.resolve_call_path(expr).and_then(|call_path| {
|
||||
if SUBSCRIPTS.contains(&call_path.as_slice()) {
|
||||
return Some(SubscriptKind::AnnotatedSubscript);
|
||||
}
|
||||
if PEP_593_SUBSCRIPTS.contains(&call_path.as_slice()) {
|
||||
return Some(SubscriptKind::PEP593AnnotatedSubscript);
|
||||
}
|
||||
|
||||
for module in &checker.settings.typing_modules {
|
||||
let module_call_path = module.split('.').collect::<Vec<_>>();
|
||||
if call_path.starts_with(&module_call_path) {
|
||||
for subscript in SUBSCRIPTS.iter() {
|
||||
if call_path.last() == subscript.last() {
|
||||
return Some(SubscriptKind::AnnotatedSubscript);
|
||||
}
|
||||
}
|
||||
}
|
||||
for module in typing_modules {
|
||||
if match_call_path(&call_path, module, member, from_imports) {
|
||||
return Some(SubscriptKind::AnnotatedSubscript);
|
||||
}
|
||||
}
|
||||
} else if let Some(modules) = PEP_593_SUBSCRIPTS.get(member) {
|
||||
for module in modules {
|
||||
if match_call_path(&call_path, module, member, from_imports)
|
||||
&& (!module.is_empty() || is_builtin(member))
|
||||
{
|
||||
return Some(SubscriptKind::PEP593AnnotatedSubscript);
|
||||
}
|
||||
}
|
||||
for module in typing_modules {
|
||||
if match_call_path(&call_path, module, member, from_imports) {
|
||||
return Some(SubscriptKind::PEP593AnnotatedSubscript);
|
||||
for subscript in PEP_593_SUBSCRIPTS.iter() {
|
||||
if call_path.last() == subscript.last() {
|
||||
return Some(SubscriptKind::PEP593AnnotatedSubscript);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
None
|
||||
|
||||
None
|
||||
})
|
||||
}
|
||||
|
||||
// See: https://peps.python.org/pep-0585/
|
||||
const PEP_585_BUILTINS_ELIGIBLE: &[(&str, &str)] = &[
|
||||
("typing", "Dict"),
|
||||
("typing", "FrozenSet"),
|
||||
("typing", "List"),
|
||||
("typing", "Set"),
|
||||
("typing", "Tuple"),
|
||||
("typing", "Type"),
|
||||
("typing_extensions", "Type"),
|
||||
const PEP_585_BUILTINS_ELIGIBLE: &[&[&str]] = &[
|
||||
&["typing", "Dict"],
|
||||
&["typing", "FrozenSet"],
|
||||
&["typing", "List"],
|
||||
&["typing", "Set"],
|
||||
&["typing", "Tuple"],
|
||||
&["typing", "Type"],
|
||||
&["typing_extensions", "Type"],
|
||||
];
|
||||
|
||||
/// Returns `true` if `Expr` represents a reference to a typing object with a
|
||||
/// PEP 585 built-in.
|
||||
pub fn is_pep585_builtin(
|
||||
expr: &Expr,
|
||||
from_imports: &FxHashMap<&str, FxHashSet<&str>>,
|
||||
import_aliases: &FxHashMap<&str, &str>,
|
||||
) -> bool {
|
||||
let call_path = dealias_call_path(collect_call_paths(expr), import_aliases);
|
||||
if !call_path.is_empty() {
|
||||
for (module, member) in PEP_585_BUILTINS_ELIGIBLE {
|
||||
if match_call_path(&call_path, module, member, from_imports) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
}
|
||||
false
|
||||
pub fn is_pep585_builtin(checker: &Checker, expr: &Expr) -> bool {
|
||||
checker.resolve_call_path(expr).map_or(false, |call_path| {
|
||||
PEP_585_BUILTINS_ELIGIBLE.contains(&call_path.as_slice())
|
||||
})
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user