Move submodule alias resolution into Context (#4543)

This commit is contained in:
Charlie Marsh 2023-05-20 12:34:10 -04:00 committed by GitHub
parent 0a3cf8ba11
commit fe7f2e2e4d
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 51 additions and 60 deletions

View file

@ -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 => {

View file

@ -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<BindingId> {
// 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,