Use the new method for autofixes across various lint rules

astral-sh · Apr 14, 2024 · 054f397 · 054f397
1 parent d2631b3
commit 054f397
Show file tree

Hide file tree

Showing 13 changed files with 160 additions and 51 deletions.
diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_bugbear/B004.py b/crates/ruff_linter/resources/test/fixtures/flake8_bugbear/B004.py
@@ -15,6 +15,16 @@ def still_a_bug():
         print("B   U   G")
 
 
+def trickier_fix_for_this_one():
+    o = object()
+
+    def callable(x):
+        return True
+
+    if hasattr(o, "__call__"):
+        print("STILL a bug!")
+
+
 def this_is_fine():
     o = object()
     if callable(o):

diff --git a/crates/ruff_linter/resources/test/fixtures/pyflakes/F901.py b/crates/ruff_linter/resources/test/fixtures/pyflakes/F901.py
@@ -4,3 +4,8 @@ def f() -> None:
 
 def g() -> None:
     raise NotImplemented
+
+
+def h() -> None:
+    NotImplementedError = "foo"
+    raise NotImplemented
diff --git a/crates/ruff_linter/src/rules/flake8_bugbear/rules/unreliable_callable_check.rs b/crates/ruff_linter/src/rules/flake8_bugbear/rules/unreliable_callable_check.rs
@@ -75,12 +75,18 @@ pub(crate) fn unreliable_callable_check(
 
     let mut diagnostic = Diagnostic::new(UnreliableCallableCheck, expr.range());
     if *function == "hasattr" {
-        if checker.semantic().is_builtin("callable") {
-            diagnostic.set_fix(Fix::safe_edit(Edit::range_replacement(
-                format!("callable({})", checker.locator().slice(obj)),
+        diagnostic.try_set_fix(|| {
+            let (import_edit, binding) = checker.importer().get_or_import_builtin_symbol(
+                "callable",
+                expr.start(),
+                checker.semantic(),
+            )?;
+            let binding_edit = Edit::range_replacement(
+                format!("{binding}({})", checker.locator().slice(obj)),
                 expr.range(),
-            )));
-        }
+            );
+            Ok(Fix::safe_edits(binding_edit, import_edit))
+        });
     }
     checker.diagnostics.push(diagnostic);
 }
diff --git a/...les/flake8_bugbear/snapshots/ruff_linter__rules__flake8_bugbear__tests__B004_B004.py.snap b/...les/flake8_bugbear/snapshots/ruff_linter__rules__flake8_bugbear__tests__B004_B004.py.snap
@@ -61,3 +61,28 @@ B004.py:14:8: B004 Using `hasattr(x, "__call__")` to test if x is callable is un
 15 |         print("B   U   G")
    |
    = help: Replace with `callable()`
+
+B004.py:24:8: B004 [*] Using `hasattr(x, "__call__")` to test if x is callable is unreliable. Use `callable(x)` for consistent results.
+   |
+22 |         return True
+23 | 
+24 |     if hasattr(o, "__call__"):
+   |        ^^^^^^^^^^^^^^^^^^^^^^ B004
+25 |         print("STILL a bug!")
+   |
+   = help: Replace with `callable()`
+
+ℹ Safe fix
+   1  |+import builtins
+1  2  | def this_is_a_bug():
+2  3  |     o = object()
+3  4  |     if hasattr(o, "__call__"):
+--------------------------------------------------------------------------------
+21 22 |     def callable(x):
+22 23 |         return True
+23 24 | 
+24    |-    if hasattr(o, "__call__"):
+   25 |+    if builtins.callable(o):
+25 26 |         print("STILL a bug!")
+26 27 | 
+27 28 |
diff --git a/crates/ruff_linter/src/rules/flake8_pyi/rules/any_eq_ne_annotation.rs b/crates/ruff_linter/src/rules/flake8_pyi/rules/any_eq_ne_annotation.rs
@@ -72,24 +72,29 @@ pub(crate) fn any_eq_ne_annotation(checker: &mut Checker, name: &str, parameters
         return;
     };
 
-    if !checker.semantic().current_scope().kind.is_class() {
+    let semantic = checker.semantic();
+
+    if !semantic.current_scope().kind.is_class() {
         return;
     }
 
-    if checker.semantic().match_typing_expr(annotation, "Any") {
+    if semantic.match_typing_expr(annotation, "Any") {
         let mut diagnostic = Diagnostic::new(
             AnyEqNeAnnotation {
                 method_name: name.to_string(),
             },
             annotation.range(),
         );
         // Ex) `def __eq__(self, obj: Any): ...`
-        if checker.semantic().is_builtin("object") {
-            diagnostic.set_fix(Fix::safe_edit(Edit::range_replacement(
-                "object".to_string(),
-                annotation.range(),
-            )));
-        }
+        diagnostic.try_set_fix(|| {
+            let (import_edit, binding) = checker.importer().get_or_import_builtin_symbol(
+                "object",
+                annotation.start(),
+                semantic,
+            )?;
+            let binding_edit = Edit::range_replacement(binding, annotation.range());
+            Ok(Fix::safe_edits(binding_edit, import_edit))
+        });
         checker.diagnostics.push(diagnostic);
     }
 }
diff --git a/crates/ruff_linter/src/rules/flake8_pyi/rules/exit_annotations.rs b/crates/ruff_linter/src/rules/flake8_pyi/rules/exit_annotations.rs
@@ -181,12 +181,15 @@ fn check_short_args_list(checker: &mut Checker, parameters: &Parameters, func_ki
                 annotation.range(),
             );
 
-            if checker.semantic().is_builtin("object") {
-                diagnostic.set_fix(Fix::safe_edit(Edit::range_replacement(
-                    "object".to_string(),
-                    annotation.range(),
-                )));
-            }
+            diagnostic.try_set_fix(|| {
+                let (import_edit, binding) = checker.importer().get_or_import_builtin_symbol(
+                    "object",
+                    annotation.start(),
+                    checker.semantic(),
+                )?;
+                let binding_edit = Edit::range_replacement(binding, annotation.range());
+                Ok(Fix::safe_edits(binding_edit, import_edit))
+            });
 
             checker.diagnostics.push(diagnostic);
         }

diff --git a/crates/ruff_linter/src/rules/pyflakes/rules/raise_not_implemented.rs b/crates/ruff_linter/src/rules/pyflakes/rules/raise_not_implemented.rs
@@ -75,11 +75,16 @@ pub(crate) fn raise_not_implemented(checker: &mut Checker, expr: &Expr) {
         return;
     };
     let mut diagnostic = Diagnostic::new(RaiseNotImplemented, expr.range());
-    if checker.semantic().is_builtin("NotImplementedError") {
-        diagnostic.set_fix(Fix::safe_edit(Edit::range_replacement(
-            "NotImplementedError".to_string(),
-            expr.range(),
-        )));
-    }
+    diagnostic.try_set_fix(|| {
+        let (import_edit, binding) = checker.importer().get_or_import_builtin_symbol(
+            "NotImplementedError",
+            expr.start(),
+            checker.semantic(),
+        )?;
+        Ok(Fix::safe_edits(
+            Edit::range_replacement(binding, expr.range()),
+            import_edit,
+        ))
+    });
     checker.diagnostics.push(diagnostic);
 }
diff --git a/...inter/src/rules/pyflakes/snapshots/ruff_linter__rules__pyflakes__tests__F901_F901.py.snap b/...inter/src/rules/pyflakes/snapshots/ruff_linter__rules__pyflakes__tests__F901_F901.py.snap
@@ -31,5 +31,27 @@ F901.py:6:11: F901 [*] `raise NotImplemented` should be `raise NotImplementedErr
 5 5 | def g() -> None:
 6   |-    raise NotImplemented
   6 |+    raise NotImplementedError
+7 7 | 
+8 8 | 
+9 9 | def h() -> None:
 
+F901.py:11:11: F901 [*] `raise NotImplemented` should be `raise NotImplementedError`
+   |
+ 9 | def h() -> None:
+10 |     NotImplementedError = "foo"
+11 |     raise NotImplemented
+   |           ^^^^^^^^^^^^^^ F901
+   |
+   = help: Use `raise NotImplementedError`
 
+ℹ Safe fix
+   1  |+import builtins
+1  2  | def f() -> None:
+2  3  |     raise NotImplemented()
+3  4  | 
+--------------------------------------------------------------------------------
+8  9  | 
+9  10 | def h() -> None:
+10 11 |     NotImplementedError = "foo"
+11    |-    raise NotImplemented
+   12 |+    raise builtins.NotImplementedError
diff --git a/crates/ruff_linter/src/rules/pyupgrade/rules/open_alias.rs b/crates/ruff_linter/src/rules/pyupgrade/rules/open_alias.rs
@@ -53,12 +53,17 @@ pub(crate) fn open_alias(checker: &mut Checker, expr: &Expr, func: &Expr) {
         .is_some_and(|qualified_name| matches!(qualified_name.segments(), ["io", "open"]))
     {
         let mut diagnostic = Diagnostic::new(OpenAlias, expr.range());
-        if checker.semantic().is_builtin("open") {
-            diagnostic.set_fix(Fix::safe_edit(Edit::range_replacement(
-                "open".to_string(),
-                func.range(),
-            )));
-        }
+        diagnostic.try_set_fix(|| {
+            let (import_edit, binding) = checker.importer().get_or_import_builtin_symbol(
+                "open",
+                expr.start(),
+                checker.semantic(),
+            )?;
+            Ok(Fix::safe_edits(
+                Edit::range_replacement(binding, func.range()),
+                import_edit,
+            ))
+        });
         checker.diagnostics.push(diagnostic);
     }
 }
diff --git a/crates/ruff_linter/src/rules/pyupgrade/rules/os_error_alias.rs b/crates/ruff_linter/src/rules/pyupgrade/rules/os_error_alias.rs
@@ -75,12 +75,17 @@ fn atom_diagnostic(checker: &mut Checker, target: &Expr) {
         },
         target.range(),
     );
-    if checker.semantic().is_builtin("OSError") {
-        diagnostic.set_fix(Fix::safe_edit(Edit::range_replacement(
-            "OSError".to_string(),
-            target.range(),
-        )));
-    }
+    diagnostic.try_set_fix(|| {
+        let (import_edit, binding) = checker.importer().get_or_import_builtin_symbol(
+            "OSError",
+            target.start(),
+            checker.semantic(),
+        )?;
+        Ok(Fix::safe_edits(
+            Edit::range_replacement(binding, target.range()),
+            import_edit,
+        ))
+    });
     checker.diagnostics.push(diagnostic);
 }
 

diff --git a/crates/ruff_linter/src/rules/pyupgrade/rules/timeout_error_alias.rs b/crates/ruff_linter/src/rules/pyupgrade/rules/timeout_error_alias.rs
@@ -89,12 +89,17 @@ fn atom_diagnostic(checker: &mut Checker, target: &Expr) {
         },
         target.range(),
     );
-    if checker.semantic().is_builtin("TimeoutError") {
-        diagnostic.set_fix(Fix::safe_edit(Edit::range_replacement(
-            "TimeoutError".to_string(),
-            target.range(),
-        )));
-    }
+    diagnostic.try_set_fix(|| {
+        let (import_edit, binding) = checker.importer().get_or_import_builtin_symbol(
+            "TimeoutError",
+            target.start(),
+            checker.semantic(),
+        )?;
+        Ok(Fix::safe_edits(
+            Edit::range_replacement(binding, target.range()),
+            import_edit,
+        ))
+    });
     checker.diagnostics.push(diagnostic);
 }
 

diff --git a/crates/ruff_linter/src/rules/pyupgrade/rules/typing_text_str_alias.rs b/crates/ruff_linter/src/rules/pyupgrade/rules/typing_text_str_alias.rs
@@ -57,12 +57,17 @@ pub(crate) fn typing_text_str_alias(checker: &mut Checker, expr: &Expr) {
         .is_some_and(|qualified_name| matches!(qualified_name.segments(), ["typing", "Text"]))
     {
         let mut diagnostic = Diagnostic::new(TypingTextStrAlias, expr.range());
-        if checker.semantic().is_builtin("str") {
-            diagnostic.set_fix(Fix::safe_edit(Edit::range_replacement(
-                "str".to_string(),
-                expr.range(),
-            )));
-        }
+        diagnostic.try_set_fix(|| {
+            let (import_edit, binding) = checker.importer().get_or_import_builtin_symbol(
+                "str",
+                expr.start(),
+                checker.semantic(),
+            )?;
+            Ok(Fix::safe_edits(
+                Edit::range_replacement(binding, expr.range()),
+                import_edit,
+            ))
+        });
         checker.diagnostics.push(diagnostic);
     }
 }
diff --git a/..._linter/src/rules/pyupgrade/snapshots/ruff_linter__rules__pyupgrade__tests__UP020.py.snap b/..._linter/src/rules/pyupgrade/snapshots/ruff_linter__rules__pyupgrade__tests__UP020.py.snap
@@ -20,7 +20,7 @@ UP020.py:3:6: UP020 [*] Use builtin `open`
 5 5 | 
 6 6 | from io import open
 
-UP020.py:8:6: UP020 Use builtin `open`
+UP020.py:8:6: UP020 [*] Use builtin `open`
   |
 6 | from io import open
 7 | 
@@ -30,4 +30,12 @@ UP020.py:8:6: UP020 Use builtin `open`
   |
   = help: Replace with builtin `open`
 
-
+ℹ Safe fix
+4  4  |     print(f.read())
+5  5  | 
+6  6  | from io import open
+   7  |+import builtins
+7  8  | 
+8     |-with open("f.txt") as f:
+   9  |+with builtins.open("f.txt") as f:
+9  10 |     print(f.read())