Handle multi-segment import-from removal (#479)

resources/test/fixtures/F401_0.py

@@ -86,3 +86,5 @@ else:
 
 
 CustomInt: TypeAlias = "np.int8 | np.int16"
+
+from foo.bar import baz

src/cst/helpers.rs

@@ -0,0 +1,42 @@
+use libcst_native::{Expression, NameOrAttribute};
+
+fn compose_call_path_inner<'a>(expr: &'a Expression, parts: &mut Vec<&'a str>) {
+    match &expr {
+        Expression::Call(expr) => {
+            compose_call_path_inner(&expr.func, parts);
+        }
+        Expression::Attribute(expr) => {
+            compose_call_path_inner(&expr.value, parts);
+            parts.push(expr.attr.value);
+        }
+        Expression::Name(expr) => {
+            parts.push(expr.value);
+        }
+        _ => {}
+    }
+}
+
+pub fn compose_call_path(expr: &Expression) -> Option<String> {
+    let mut segments = vec![];
+    compose_call_path_inner(expr, &mut segments);
+    if segments.is_empty() {
+        None
+    } else {
+        Some(segments.join("."))
+    }
+}
+
+pub fn compose_module_path(module: &NameOrAttribute) -> String {
+    match module {
+        NameOrAttribute::N(name) => name.value.to_string(),
+        NameOrAttribute::A(attr) => {
+            let name = attr.attr.value;
+            let prefix = compose_call_path(&attr.value);
+            if let Some(prefix) = prefix {
+                format!("{prefix}.{name}")
+            } else {
+                name.to_string()
+            }
+        }
+    }
+}

src/cst/mod.rs

@@ -0,0 +1 @@
+pub mod helpers;

src/lib.rs

@@ -19,6 +19,7 @@ mod check_lines;
 pub mod checks;
 pub mod cli;
 pub mod code_gen;
+mod cst;
 mod docstrings;
 mod flake8_bugbear;
 mod flake8_builtins;

src/pyflakes/fixes.rs

@@ -1,11 +1,11 @@
-use libcst_native::ImportNames::Aliases;
-use libcst_native::NameOrAttribute::N;
-use libcst_native::{Codegen, SmallStatement, Statement};
+use anyhow::Result;
+use libcst_native::{Codegen, ImportNames, NameOrAttribute, SmallStatement, Statement};
 use rustpython_ast::Stmt;
 
 use crate::ast::operations::SourceCodeLocator;
 use crate::ast::types::Range;
 use crate::autofix::{helpers, Fix};
+use crate::cst::helpers::compose_module_path;
 
 /// Generate a Fix to remove any unused imports from an `import` statement.
 pub fn remove_unused_imports(
@@ -14,7 +14,7 @@ pub fn remove_unused_imports(
     stmt: &Stmt,
     parent: Option<&Stmt>,
     deleted: &[&Stmt],
-) -> anyhow::Result<Fix> {
+) -> Result<Fix> {
     let mut tree = match libcst_native::parse_module(
         locator.slice_source_code_range(&Range::from_located(stmt)),
         None,
@@ -43,7 +43,7 @@ pub fn remove_unused_imports(
     // Identify unused imports from within the `import from`.
     let mut removable = vec![];
     for (index, alias) in aliases.iter().enumerate() {
-        if let N(import_name) = &alias.name {
+        if let NameOrAttribute::N(import_name) = &alias.name {
             if full_names.contains(&import_name.value) {
                 removable.push(index);
             }
@@ -79,7 +79,7 @@ pub fn remove_unused_import_froms(
     stmt: &Stmt,
     parent: Option<&Stmt>,
     deleted: &[&Stmt],
-) -> anyhow::Result<Fix> {
+) -> Result<Fix> {
     let mut tree = match libcst_native::parse_module(
         locator.slice_source_code_range(&Range::from_located(stmt)),
         None,
@@ -100,7 +100,8 @@ pub fn remove_unused_import_froms(
             "Expected node to be: SmallStatement::ImportFrom."
         ));
     };
-    let aliases = if let Aliases(aliases) = &mut body.names {
+
+    let aliases = if let ImportNames::Aliases(aliases) = &mut body.names {
         aliases
     } else {
         return Err(anyhow::anyhow!("Expected node to be: Aliases."));
@@ -112,13 +113,16 @@ pub fn remove_unused_import_froms(
     // Identify unused imports from within the `import from`.
     let mut removable = vec![];
     for (index, alias) in aliases.iter().enumerate() {
-        if let N(name) = &alias.name {
-            let import_name = if let Some(N(module_name)) = &body.module {
-                format!("{}.{}", module_name.value, name.value)
-            } else {
-                name.value.to_string()
-            };
-            if full_names.contains(&import_name.as_str()) {
+        if let NameOrAttribute::N(name) = &alias.name {
+            let import_name = name.value.to_string();
+            let full_name = body
+                .module
+                .as_ref()
+                .map(compose_module_path)
+                .map(|module_name| format!("{module_name}.{import_name}"))
+                .unwrap_or(import_name);
+
+            if full_names.contains(&full_name.as_str()) {
                 removable.push(index);
             }
         }

src/snapshots/ruff__linter__tests__F401_F401_0.py.snap

@@ -135,4 +135,23 @@ expression: checks
         row: 53
         column: 22
     applied: false
+- kind:
+    UnusedImport:
+      - foo.bar.baz
+  location:
+    row: 90
+    column: 1
+  end_location:
+    row: 90
+    column: 24
+  fix:
+    patch:
+      content: ""
+      location:
+        row: 90
+        column: 1
+      end_location:
+        row: 91
+        column: 1
+    applied: false