Mark aliased submodule imports as used (#374)

resources/test/fixtures/F401.py

@@ -65,3 +65,9 @@ b = Union["Nut", None]
 c = cast("Vegetable", b)
 
 Field = lambda default=MISSING: field(default=default)
+
+
+import pyarrow as pa
+import pyarrow.csv
+
+print(pa.csv.read_csv("test.csv"))

src/ast/types.rs

@@ -74,9 +74,9 @@ pub enum BindingKind {
     Export(Vec<String>),
     FutureImportation,
     StarImportation,
-    Importation(String, BindingContext),
+    Importation(String, String, BindingContext),
-    FromImportation(String, BindingContext),
+    FromImportation(String, String, BindingContext),
-    SubmoduleImportation(String, BindingContext),
+    SubmoduleImportation(String, String, BindingContext),
 }
 
 #[derive(Clone, Debug)]

src/check_ast.rs

@@ -364,13 +364,16 @@ where
 
                 for alias in names {
                     if alias.node.name.contains('.') && alias.node.asname.is_none() {
-                        // TODO(charlie): Multiple submodule imports with the same parent module
+                        // Given `import foo.bar`, `name` would be "foo", and `full_name` would be
-                        // will be merged into a single binding.
+                        // "foo.bar".
+                        let name = alias.node.name.split('.').next().unwrap();
+                        let full_name = &alias.node.name;
                         self.add_binding(
-                            alias.node.name.split('.').next().unwrap().to_string(),
+                            name.to_string(),
                             Binding {
                                 kind: BindingKind::SubmoduleImportation(
-                                    alias.node.name.to_string(),
+                                    name.to_string(),
+                                    full_name.to_string(),
                                     self.binding_context(),
                                 ),
                                 used: None,
@@ -382,19 +385,17 @@ where
                             self.check_builtin_shadowing(asname, Range::from_located(stmt), false);
                         }
 
+                        // Given `import foo`, `name` and `full_name` would both be `foo`.
+                        // Given `import foo as bar`, `name` would be `bar` and `full_name` would
+                        // be `foo`.
+                        let name = alias.node.asname.as_ref().unwrap_or(&alias.node.name);
+                        let full_name = &alias.node.name;
                         self.add_binding(
-                            alias
+                            name.to_string(),
-                                .node
-                                .asname
-                                .clone()
-                                .unwrap_or_else(|| alias.node.name.clone()),
                             Binding {
                                 kind: BindingKind::Importation(
-                                    alias
+                                    name.to_string(),
-                                        .node
+                                    full_name.to_string(),
-                                        .asname
-                                        .clone()
-                                        .unwrap_or_else(|| alias.node.name.clone()),
                                     self.binding_context(),
                                 ),
                                 used: None,
@@ -423,14 +424,10 @@ where
                 }
 
                 for alias in names {
-                    let name = alias
-                        .node
-                        .asname
-                        .clone()
-                        .unwrap_or_else(|| alias.node.name.clone());
                     if let Some("__future__") = module.as_deref() {
+                        let name = alias.node.asname.as_ref().unwrap_or(&alias.node.name);
                         self.add_binding(
-                            name,
+                            name.to_string(),
                             Binding {
                                 kind: BindingKind::FutureImportation,
                                 used: Some((
@@ -509,18 +506,26 @@ where
                             self.check_builtin_shadowing(asname, Range::from_located(stmt), false);
                         }
 
-                        let binding = Binding {
+                        // Given `from foo import bar`, `name` would be "bar" and `full_name` would
-                            kind: BindingKind::FromImportation(
+                        // be "foo.bar". Given `from foo import bar as baz`, `name` would be "baz"
-                                match module {
+                        // and `full_name` would be "foo.bar".
-                                    None => name.clone(),
+                        let name = alias.node.asname.as_ref().unwrap_or(&alias.node.name);
-                                    Some(parent) => format!("{}.{}", parent, name),
+                        let full_name = match module {
-                                },
+                            None => alias.node.name.to_string(),
-                                self.binding_context(),
+                            Some(parent) => format!("{}.{}", parent, alias.node.name),
-                            ),
-                            used: None,
-                            range: Range::from_located(stmt),
                         };
-                        self.add_binding(name, binding)
+                        self.add_binding(
+                            name.to_string(),
+                            Binding {
+                                kind: BindingKind::FromImportation(
+                                    name.to_string(),
+                                    full_name,
+                                    self.binding_context(),
+                                ),
+                                used: None,
+                                range: Range::from_located(stmt),
+                            },
+                        )
                     }
                 }
             }
@@ -1255,6 +1260,47 @@ impl CheckLocator for Checker<'_> {
     }
 }
 
+fn try_mark_used(scope: &mut Scope, scope_id: usize, id: &str, expr: &Expr) -> bool {
+    let alias = if let Some(binding) = scope.values.get_mut(id) {
+        // Mark the binding as used.
+        binding.used = Some((scope_id, Range::from_located(expr)));
+
+        // 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.
+        //
+        // This handles code like:
+        //   import pyarrow as pa
+        //   import pyarrow.csv
+        //   print(pa.csv.read_csv("test.csv"))
+        if let BindingKind::Importation(name, full_name, _)
+        | BindingKind::FromImportation(name, full_name, _)
+        | BindingKind::SubmoduleImportation(name, full_name, _) = &binding.kind
+        {
+            let has_alias = full_name
+                .split('.')
+                .last()
+                .map(|segment| segment != name)
+                .unwrap_or_default();
+            if has_alias {
+                // Clone the alias. (We'll mutate it below.)
+                full_name.to_string()
+            } else {
+                return true;
+            }
+        } else {
+            return true;
+        }
+    } else {
+        return false;
+    };
+
+    // Mark the sub-importation as used.
+    if let Some(binding) = scope.values.get_mut(&alias) {
+        binding.used = Some((scope_id, Range::from_located(expr)));
+    }
+    true
+}
+
 impl<'a> Checker<'a> {
     pub fn add_check(&mut self, check: Check) {
         self.checks.push(check);
@@ -1338,7 +1384,11 @@ impl<'a> Checker<'a> {
                     && matches!(binding.kind, BindingKind::LoopVar)
                     && matches!(
                         existing.kind,
-                        BindingKind::Importation(_, _) | BindingKind::FromImportation(_, _)
+                        BindingKind::Importation(_, _, _)
+                            | BindingKind::FromImportation(_, _, _)
+                            | BindingKind::SubmoduleImportation(_, _, _)
+                            | BindingKind::StarImportation
+                            | BindingKind::FutureImportation
                     )
                 {
                     self.checks.push(Check::new(
@@ -1377,8 +1427,8 @@ impl<'a> Checker<'a> {
                         continue;
                     }
                 }
-                if let Some(binding) = scope.values.get_mut(id) {
+
-                    binding.used = Some((scope_id, Range::from_located(expr)));
+                if try_mark_used(scope, scope_id, id, expr) {
                     return;
                 }
 
@@ -1713,7 +1763,7 @@ impl<'a> Checker<'a> {
 
                     if !used {
                         match &binding.kind {
-                            BindingKind::FromImportation(full_name, context) => {
+                            BindingKind::FromImportation(_, full_name, context) => {
                                 let full_names = unused
                                     .entry((
                                         ImportKind::ImportFrom,
@@ -1723,8 +1773,8 @@ impl<'a> Checker<'a> {
                                     .or_default();
                                 full_names.push(full_name);
                             }
-                            BindingKind::Importation(full_name, context)
+                            BindingKind::Importation(_, full_name, context)
-                            | BindingKind::SubmoduleImportation(full_name, context) => {
+                            | BindingKind::SubmoduleImportation(_, full_name, context) => {
                                 let full_names = unused
                                     .entry((
                                         ImportKind::Import,