[RUF007] - add autofix (zip-instead-of-pairwise)

astral-sh · Aug 5, 2024 · 533e022 · 533e022
1 parent 7a2c75e
commit 533e022
Show file tree

Hide file tree

Showing 5 changed files with 313 additions and 11 deletions.
diff --git a/crates/ruff_linter/src/checkers/ast/analyze/expression.rs b/crates/ruff_linter/src/checkers/ast/analyze/expression.rs
@@ -852,7 +852,7 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
             }
             if checker.enabled(Rule::ZipInsteadOfPairwise) {
                 if checker.settings.target_version >= PythonVersion::Py310 {
-                    ruff::rules::zip_instead_of_pairwise(checker, func, args);
+                    ruff::rules::zip_instead_of_pairwise(checker, call);
                 }
             }
             if checker.any_enabled(&[

diff --git a/crates/ruff_linter/src/rules/ruff/mod.rs b/crates/ruff_linter/src/rules/ruff/mod.rs
@@ -346,4 +346,22 @@ mod tests {
         assert_messages!(snapshot, messages);
         Ok(())
     }
+
+    #[test_case(Rule::ZipInsteadOfPairwise, Path::new("RUF007.py"))]
+    fn preview_rules(rule_code: Rule, path: &Path) -> Result<()> {
+        let snapshot = format!(
+            "preview__{}_{}",
+            rule_code.noqa_code(),
+            path.to_string_lossy()
+        );
+        let diagnostics = test_path(
+            Path::new("ruff").join(path).as_path(),
+            &settings::LinterSettings {
+                preview: PreviewMode::Enabled,
+                ..settings::LinterSettings::for_rule(rule_code)
+            },
+        )?;
+        assert_messages!(snapshot, diagnostics);
+        Ok(())
+    }
 }
diff --git a/crates/ruff_linter/src/rules/ruff/rules/zip_instead_of_pairwise.rs b/crates/ruff_linter/src/rules/ruff/rules/zip_instead_of_pairwise.rs
@@ -1,9 +1,9 @@
-use ruff_diagnostics::{Diagnostic, Violation};
+use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation};
 use ruff_macros::{derive_message_formats, violation};
-use ruff_python_ast::{self as ast, Expr, Int};
+use ruff_python_ast::{self as ast, Arguments, Expr, Int};
 use ruff_text_size::Ranged;
 
-use crate::checkers::ast::Checker;
+use crate::{checkers::ast::Checker, importer::ImportRequest};
 
 /// ## What it does
 /// Checks for use of `zip()` to iterate over successive pairs of elements.
@@ -35,10 +35,15 @@ use crate::checkers::ast::Checker;
 pub struct ZipInsteadOfPairwise;
 
 impl Violation for ZipInsteadOfPairwise {
+    const FIX_AVAILABILITY: FixAvailability = FixAvailability::Sometimes;
     #[derive_message_formats]
     fn message(&self) -> String {
         format!("Prefer `itertools.pairwise()` over `zip()` when iterating over successive pairs")
     }
+
+    fn fix_title(&self) -> Option<String> {
+        Some("Replace `zip()` with `itertools.pairwise()`".to_string())
+    }
 }
 
 #[derive(Debug)]
@@ -95,9 +100,15 @@ fn match_slice_info(expr: &Expr) -> Option<SliceInfo> {
 }
 
 /// RUF007
-pub(crate) fn zip_instead_of_pairwise(checker: &mut Checker, func: &Expr, args: &[Expr]) {
+pub(crate) fn zip_instead_of_pairwise(checker: &mut Checker, call: &ast::ExprCall) {
+    let ast::ExprCall {
+        func,
+        arguments: Arguments { args, .. },
+        ..
+    } = call;
+
     // Require exactly two positional arguments.
-    let [first, second] = args else {
+    let [first, second] = args.as_ref() else {
         return;
     };
 
@@ -139,7 +150,20 @@ pub(crate) fn zip_instead_of_pairwise(checker: &mut Checker, func: &Expr, args:
         return;
     }
 
-    checker
-        .diagnostics
-        .push(Diagnostic::new(ZipInsteadOfPairwise, func.range()));
+    let mut diagnostic = Diagnostic::new(ZipInsteadOfPairwise, func.range());
+
+    if checker.settings.preview.is_enabled() {
+        diagnostic.try_set_fix(|| {
+            let (import_edit, binding) = checker.importer().get_or_import_symbol(
+                &ImportRequest::import("itertools", "pairwise"),
+                func.start(),
+                checker.semantic(),
+            )?;
+            let reference_edit =
+                Edit::range_replacement(format!("{binding}({})", first_arg_info.id), call.range());
+            Ok(Fix::unsafe_edits(import_edit, [reference_edit]))
+        });
+    }
+
+    checker.diagnostics.push(diagnostic);
 }
diff --git a/...ff_linter/src/rules/ruff/snapshots/ruff_linter__rules__ruff__tests__RUF007_RUF007.py.snap b/...ff_linter/src/rules/ruff/snapshots/ruff_linter__rules__ruff__tests__RUF007_RUF007.py.snap
@@ -9,6 +9,7 @@ RUF007.py:16:1: RUF007 Prefer `itertools.pairwise()` over `zip()` when iterating
 17 | zip(input, input[1::1])
 18 | zip(input[:-1], input[1:])
    |
+   = help: Replace `zip()` with `itertools.pairwise()`
 
 RUF007.py:17:1: RUF007 Prefer `itertools.pairwise()` over `zip()` when iterating over successive pairs
    |
@@ -19,6 +20,7 @@ RUF007.py:17:1: RUF007 Prefer `itertools.pairwise()` over `zip()` when iterating
 18 | zip(input[:-1], input[1:])
 19 | zip(input[1:], input[2:])
    |
+   = help: Replace `zip()` with `itertools.pairwise()`
 
 RUF007.py:18:1: RUF007 Prefer `itertools.pairwise()` over `zip()` when iterating over successive pairs
    |
@@ -29,6 +31,7 @@ RUF007.py:18:1: RUF007 Prefer `itertools.pairwise()` over `zip()` when iterating
 19 | zip(input[1:], input[2:])
 20 | zip(input[1:-1], input[2:])
    |
+   = help: Replace `zip()` with `itertools.pairwise()`
 
 RUF007.py:19:1: RUF007 Prefer `itertools.pairwise()` over `zip()` when iterating over successive pairs
    |
@@ -39,6 +42,7 @@ RUF007.py:19:1: RUF007 Prefer `itertools.pairwise()` over `zip()` when iterating
 20 | zip(input[1:-1], input[2:])
 21 | list(zip(input, input[1:]))
    |
+   = help: Replace `zip()` with `itertools.pairwise()`
 
 RUF007.py:20:1: RUF007 Prefer `itertools.pairwise()` over `zip()` when iterating over successive pairs
    |
@@ -49,6 +53,7 @@ RUF007.py:20:1: RUF007 Prefer `itertools.pairwise()` over `zip()` when iterating
 21 | list(zip(input, input[1:]))
 22 | list(zip(input[:-1], input[1:]))
    |
+   = help: Replace `zip()` with `itertools.pairwise()`
 
 RUF007.py:21:6: RUF007 Prefer `itertools.pairwise()` over `zip()` when iterating over successive pairs
    |
@@ -59,6 +64,7 @@ RUF007.py:21:6: RUF007 Prefer `itertools.pairwise()` over `zip()` when iterating
 22 | list(zip(input[:-1], input[1:]))
 23 | zip(foo[:-1], foo[1:], strict=True)
    |
+   = help: Replace `zip()` with `itertools.pairwise()`
 
 RUF007.py:22:6: RUF007 Prefer `itertools.pairwise()` over `zip()` when iterating over successive pairs
    |
@@ -69,6 +75,7 @@ RUF007.py:22:6: RUF007 Prefer `itertools.pairwise()` over `zip()` when iterating
 23 | zip(foo[:-1], foo[1:], strict=True)
 24 | zip(foo[:-1], foo[1:], strict=False)
    |
+   = help: Replace `zip()` with `itertools.pairwise()`
 
 RUF007.py:23:1: RUF007 Prefer `itertools.pairwise()` over `zip()` when iterating over successive pairs
    |
@@ -79,6 +86,7 @@ RUF007.py:23:1: RUF007 Prefer `itertools.pairwise()` over `zip()` when iterating
 24 | zip(foo[:-1], foo[1:], strict=False)
 25 | zip(foo[:-1], foo[1:], strict=bool(foo))
    |
+   = help: Replace `zip()` with `itertools.pairwise()`
 
 RUF007.py:24:1: RUF007 Prefer `itertools.pairwise()` over `zip()` when iterating over successive pairs
    |
@@ -88,6 +96,7 @@ RUF007.py:24:1: RUF007 Prefer `itertools.pairwise()` over `zip()` when iterating
    | ^^^ RUF007
 25 | zip(foo[:-1], foo[1:], strict=bool(foo))
    |
+   = help: Replace `zip()` with `itertools.pairwise()`
 
 RUF007.py:25:1: RUF007 Prefer `itertools.pairwise()` over `zip()` when iterating over successive pairs
    |
@@ -96,5 +105,4 @@ RUF007.py:25:1: RUF007 Prefer `itertools.pairwise()` over `zip()` when iterating
 25 | zip(foo[:-1], foo[1:], strict=bool(foo))
    | ^^^ RUF007
    |
-
-
+   = help: Replace `zip()` with `itertools.pairwise()`