Move submodule alias resolution into Context (#4543)

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) => {
+            ResolvedReference::Resolved(..) | ResolvedReference::ImplicitGlobal => {
-                // 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 => {
                 // Nothing to do.
             }
             ResolvedReference::StarImport => {

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<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,