From 16c2e3a99530497dda9e044b295e07193a8c1266 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Wed, 26 Oct 2022 16:36:12 -0400 Subject: [PATCH] Handle multi-segment import-from removal (#479) --- resources/test/fixtures/F401_0.py | 2 + src/{ast.rs => ast/mod.rs} | 0 src/cst/helpers.rs | 42 +++++++++++++++++++ src/cst/mod.rs | 1 + src/lib.rs | 1 + src/pyflakes/fixes.rs | 32 +++++++------- .../ruff__linter__tests__F401_F401_0.py.snap | 19 +++++++++ 7 files changed, 83 insertions(+), 14 deletions(-) rename src/{ast.rs => ast/mod.rs} (100%) create mode 100644 src/cst/helpers.rs create mode 100644 src/cst/mod.rs diff --git a/resources/test/fixtures/F401_0.py b/resources/test/fixtures/F401_0.py index 24afc4c6fe..d9e534fa6d 100644 --- a/resources/test/fixtures/F401_0.py +++ b/resources/test/fixtures/F401_0.py @@ -86,3 +86,5 @@ else: CustomInt: TypeAlias = "np.int8 | np.int16" + +from foo.bar import baz diff --git a/src/ast.rs b/src/ast/mod.rs similarity index 100% rename from src/ast.rs rename to src/ast/mod.rs diff --git a/src/cst/helpers.rs b/src/cst/helpers.rs new file mode 100644 index 0000000000..942ee625e1 --- /dev/null +++ b/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 { + 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() + } + } + } +} diff --git a/src/cst/mod.rs b/src/cst/mod.rs new file mode 100644 index 0000000000..1630fabcd1 --- /dev/null +++ b/src/cst/mod.rs @@ -0,0 +1 @@ +pub mod helpers; diff --git a/src/lib.rs b/src/lib.rs index ee793a0356..b3da0d8e1f 100644 --- a/src/lib.rs +++ b/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; diff --git a/src/pyflakes/fixes.rs b/src/pyflakes/fixes.rs index f8fd68ec3b..b92043f1cc 100644 --- a/src/pyflakes/fixes.rs +++ b/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 { +) -> Result { 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 { +) -> Result { 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); } } diff --git a/src/snapshots/ruff__linter__tests__F401_F401_0.py.snap b/src/snapshots/ruff__linter__tests__F401_F401_0.py.snap index df3c4407a5..9f8c5c9c11 100644 --- a/src/snapshots/ruff__linter__tests__F401_F401_0.py.snap +++ b/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