From fe7f2e2e4db2e709caa80d59ba970d242d9cf560 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Sat, 20 May 2023 12:34:10 -0400 Subject: [PATCH] Move submodule alias resolution into `Context` (#4543) --- crates/ruff/src/checkers/ast/mod.rs | 58 +--------------------- crates/ruff_python_semantic/src/context.rs | 53 ++++++++++++++++++-- 2 files changed, 51 insertions(+), 60 deletions(-) diff --git a/crates/ruff/src/checkers/ast/mod.rs b/crates/ruff/src/checkers/ast/mod.rs index 033fa6b755..2ed17be534 100644 --- a/crates/ruff/src/checkers/ast/mod.rs +++ b/crates/ruff/src/checkers/ast/mod.rs @@ -4747,63 +4747,7 @@ impl<'a> Checker<'a> { return; }; match self.ctx.resolve_reference(id, expr.range()) { - ResolvedReference::Resolved(scope_id, binding_id) => { - // If the name of the sub-importation is the same as an alias of another - // importation and the alias is used, that sub-importation should be - // marked as used too. - // - // For example, mark `pa` as used in: - // - // ```python - // import pyarrow as pa - // import pyarrow.csv - // print(pa.csv.read_csv("test.csv")) - // ``` - match &self.ctx.bindings[binding_id].kind { - BindingKind::Importation(Importation { name, full_name }) - | BindingKind::SubmoduleImportation(SubmoduleImportation { name, full_name }) => - { - let has_alias = full_name - .split('.') - .last() - .map(|segment| &segment != name) - .unwrap_or_default(); - if has_alias { - // Mark the sub-importation as used. - if let Some(binding_id) = self.ctx.scopes[scope_id].get(full_name) { - let context = self.ctx.execution_context(); - self.ctx.bindings[*binding_id].mark_used( - self.ctx.scope_id, - expr.range(), - context, - ); - } - } - } - BindingKind::FromImportation(FromImportation { name, full_name }) => { - let has_alias = full_name - .split('.') - .last() - .map(|segment| &segment != name) - .unwrap_or_default(); - if has_alias { - // Mark the sub-importation as used. - if let Some(binding_id) = - self.ctx.scopes[scope_id].get(full_name.as_str()) - { - let context = self.ctx.execution_context(); - self.ctx.bindings[*binding_id].mark_used( - self.ctx.scope_id, - expr.range(), - context, - ); - } - } - } - _ => {} - } - } - ResolvedReference::ImplicitGlobal => { + ResolvedReference::Resolved(..) | ResolvedReference::ImplicitGlobal => { // Nothing to do. } ResolvedReference::StarImport => { diff --git a/crates/ruff_python_semantic/src/context.rs b/crates/ruff_python_semantic/src/context.rs index 3f3f91c7ae..1831b991c5 100644 --- a/crates/ruff_python_semantic/src/context.rs +++ b/crates/ruff_python_semantic/src/context.rs @@ -120,9 +120,16 @@ impl<'a> Context<'a> { // should prefer it over local resolutions. if self.in_deferred_type_definition() { if let Some(binding_id) = self.scopes.global().get(symbol) { + // Mark the binding as used. let context = self.execution_context(); self.bindings[*binding_id].mark_used(ScopeId::global(), range, context); - return ResolvedReference::Resolved(ScopeId::global(), *binding_id); + + // Mark any submodule aliases as used. + if let Some(binding_id) = self.resolve_submodule(ScopeId::global(), *binding_id) { + self.bindings[binding_id].mark_used(ScopeId::global(), range, context); + } + + return ResolvedReference::Resolved(*binding_id); } } @@ -151,6 +158,11 @@ impl<'a> Context<'a> { let context = self.execution_context(); self.bindings[*binding_id].mark_used(self.scope_id, range, context); + // Mark any submodule aliases as used. + if let Some(binding_id) = self.resolve_submodule(scope_id, *binding_id) { + self.bindings[binding_id].mark_used(ScopeId::global(), range, context); + } + // But if it's a type annotation, don't treat it as resolved, unless we're in a // forward reference. For example, given: // @@ -167,7 +179,7 @@ impl<'a> Context<'a> { continue; } - return ResolvedReference::Resolved(scope_id, *binding_id); + return ResolvedReference::Resolved(*binding_id); } // Allow usages of `__module__` and `__qualname__` within class scopes, e.g.: @@ -202,6 +214,41 @@ impl<'a> Context<'a> { } } + /// Given a `BindingId`, return the `BindingId` of the submodule import that it aliases. + fn resolve_submodule(&self, scope_id: ScopeId, binding_id: BindingId) -> Option { + // If the name of a submodule import is the same as an alias of another import, and the + // alias is used, then the submodule import should be marked as used too. + // + // For example, mark `pyarrow.csv` as used in: + // + // ```python + // import pyarrow as pa + // import pyarrow.csv + // print(pa.csv.read_csv("test.csv")) + // ``` + let (name, full_name) = match &self.bindings[binding_id].kind { + BindingKind::Importation(Importation { name, full_name }) => (*name, *full_name), + BindingKind::SubmoduleImportation(SubmoduleImportation { name, full_name }) => { + (*name, *full_name) + } + BindingKind::FromImportation(FromImportation { name, full_name }) => { + (*name, full_name.as_str()) + } + _ => return None, + }; + + let has_alias = full_name + .split('.') + .last() + .map(|segment| segment != name) + .unwrap_or_default(); + if !has_alias { + return None; + } + + self.scopes[scope_id].get(full_name).copied() + } + /// Resolves the [`Expr`] to a fully-qualified symbol-name, if `value` resolves to an imported /// or builtin symbol. /// @@ -801,7 +848,7 @@ pub struct Snapshot { #[derive(Debug)] pub enum ResolvedReference { /// The reference is resolved to a specific binding. - Resolved(ScopeId, BindingId), + Resolved(BindingId), /// The reference is resolved to a context-specific, implicit global (e.g., `__class__` within /// a class scope). ImplicitGlobal,