Add fix to automatically remove print and pprint statements (#9208)

charliermarsh · web-flow · commit 07b293d94922 · 2023-12-20T05:35:30.000Z
Closes #9207.
diff --git a/crates/ruff_linter/src/rules/flake8_print/rules/print_call.rs b/crates/ruff_linter/src/rules/flake8_print/rules/print_call.rs
@@ -1,10 +1,10 @@
-use ruff_diagnostics::{Diagnostic, Violation};
+use ruff_diagnostics::{Diagnostic, Fix, FixAvailability, Violation};
 use ruff_macros::{derive_message_formats, violation};
-
-use ruff_python_ast::{self as ast};
+use ruff_python_ast as ast;
 use ruff_text_size::Ranged;
 
 use crate::checkers::ast::Checker;
+use crate::fix::edits::delete_stmt;
 use crate::registry::AsRule;
 
 /// ## What it does
@@ -28,14 +28,24 @@ use crate::registry::AsRule;
 /// def add_numbers(a, b):
 ///     return a + b
 /// ```
+///
+/// ## Fix safety
+/// This rule's fix is marked as unsafe, as it may remove `print` statements
+/// that are used beyond debugging purposes.
 #[violation]
 pub struct Print;
 
 impl Violation for Print {
+    const FIX_AVAILABILITY: FixAvailability = FixAvailability::Sometimes;
+
     #[derive_message_formats]
     fn message(&self) -> String {
         format!("`print` found")
     }
+
+    fn fix_title(&self) -> Option<String> {
+        Some("Remove `print`".to_string())
+    }
 }
 
 /// ## What it does
@@ -65,19 +75,29 @@ impl Violation for Print {
 ///     dict_c = {**dict_a, **dict_b}
 ///     return dict_c
 /// ```
+///
+/// ## Fix safety
+/// This rule's fix is marked as unsafe, as it may remove `pprint` statements
+/// that are used beyond debugging purposes.
 #[violation]
 pub struct PPrint;
 
 impl Violation for PPrint {
+    const FIX_AVAILABILITY: FixAvailability = FixAvailability::Sometimes;
+
     #[derive_message_formats]
     fn message(&self) -> String {
         format!("`pprint` found")
     }
+
+    fn fix_title(&self) -> Option<String> {
+        Some("Remove `pprint`".to_string())
+    }
 }
 
 /// T201, T203
 pub(crate) fn print_call(checker: &mut Checker, call: &ast::ExprCall) {
-    let diagnostic = {
+    let mut diagnostic = {
         let call_path = checker.semantic().resolve_call_path(&call.func);
         if call_path
             .as_ref()
@@ -113,5 +133,15 @@ pub(crate) fn print_call(checker: &mut Checker, call: &ast::ExprCall) {
         return;
     }
 
+    // Remove the `print`, if it's a standalone statement.
+    if checker.semantic().current_expression_parent().is_none() {
+        let statement = checker.semantic().current_statement();
+        let parent = checker.semantic().current_statement_parent();
+        let edit = delete_stmt(statement, parent, checker.locator(), checker.indexer());
+        diagnostic.set_fix(Fix::unsafe_edit(edit).isolate(Checker::isolation(
+            checker.semantic().current_statement_parent_id(),
+        )));
+    }
+
     checker.diagnostics.push(diagnostic);
 }
diff --git a/crates/ruff_linter/src/rules/flake8_print/snapshots/ruff_linter__rules__flake8_print__tests__T201_T201.py.snap b/crates/ruff_linter/src/rules/flake8_print/snapshots/ruff_linter__rules__flake8_print__tests__T201_T201.py.snap
@@ -1,7 +1,7 @@
 ---
 source: crates/ruff_linter/src/rules/flake8_print/mod.rs
 ---
-T201.py:4:1: T201 `print` found
+T201.py:4:1: T201 [*] `print` found
   |
 2 | import tempfile
 3 | 
@@ -10,26 +10,56 @@ T201.py:4:1: T201 `print` found
 5 | print("Hello, world!", file=None)  # T201
 6 | print("Hello, world!", file=sys.stdout)  # T201
   |
+  = help: Remove `print`
 
-T201.py:5:1: T201 `print` found
+ℹ Unsafe fix
+1 1 | import sys
+2 2 | import tempfile
+3 3 | 
+4   |-print("Hello, world!")  # T201
+5 4 | print("Hello, world!", file=None)  # T201
+6 5 | print("Hello, world!", file=sys.stdout)  # T201
+7 6 | print("Hello, world!", file=sys.stderr)  # T201
+
+T201.py:5:1: T201 [*] `print` found
   |
 4 | print("Hello, world!")  # T201
 5 | print("Hello, world!", file=None)  # T201
   | ^^^^^ T201
 6 | print("Hello, world!", file=sys.stdout)  # T201
 7 | print("Hello, world!", file=sys.stderr)  # T201
   |
+  = help: Remove `print`
+
+ℹ Unsafe fix
+2 2 | import tempfile
+3 3 | 
+4 4 | print("Hello, world!")  # T201
+5   |-print("Hello, world!", file=None)  # T201
+6 5 | print("Hello, world!", file=sys.stdout)  # T201
+7 6 | print("Hello, world!", file=sys.stderr)  # T201
+8 7 | 
 
-T201.py:6:1: T201 `print` found
+T201.py:6:1: T201 [*] `print` found
   |
 4 | print("Hello, world!")  # T201
 5 | print("Hello, world!", file=None)  # T201
 6 | print("Hello, world!", file=sys.stdout)  # T201
   | ^^^^^ T201
 7 | print("Hello, world!", file=sys.stderr)  # T201
   |
+  = help: Remove `print`
 
-T201.py:7:1: T201 `print` found
+ℹ Unsafe fix
+3 3 | 
+4 4 | print("Hello, world!")  # T201
+5 5 | print("Hello, world!", file=None)  # T201
+6   |-print("Hello, world!", file=sys.stdout)  # T201
+7 6 | print("Hello, world!", file=sys.stderr)  # T201
+8 7 | 
+9 8 | with tempfile.NamedTemporaryFile() as fp:
+
+T201.py:7:1: T201 [*] `print` found
   |
 5 | print("Hello, world!", file=None)  # T201
 6 | print("Hello, world!", file=sys.stdout)  # T201
@@ -38,5 +68,15 @@ T201.py:7:1: T201 `print` found
 8 | 
 9 | with tempfile.NamedTemporaryFile() as fp:
   |
+  = help: Remove `print`
+
+ℹ Unsafe fix
+4 4 | print("Hello, world!")  # T201
+5 5 | print("Hello, world!", file=None)  # T201
+6 6 | print("Hello, world!", file=sys.stdout)  # T201
+7   |-print("Hello, world!", file=sys.stderr)  # T201
+8 7 | 
+9 8 | with tempfile.NamedTemporaryFile() as fp:
+10 9 |     print("Hello, world!", file=fp)  # OK
 
 
diff --git a/crates/ruff_linter/src/rules/flake8_print/snapshots/ruff_linter__rules__flake8_print__tests__T203_T203.py.snap b/crates/ruff_linter/src/rules/flake8_print/snapshots/ruff_linter__rules__flake8_print__tests__T203_T203.py.snap
@@ -1,7 +1,7 @@
 ---
 source: crates/ruff_linter/src/rules/flake8_print/mod.rs
 ---
-T203.py:3:1: T203 `pprint` found
+T203.py:3:1: T203 [*] `pprint` found
   |
 1 | from pprint import pprint
 2 | 
@@ -10,8 +10,17 @@ T203.py:3:1: T203 `pprint` found
 4 | 
 5 | import pprint
   |
+  = help: Remove `pprint`
 
-T203.py:7:1: T203 `pprint` found
+ℹ Unsafe fix
+1 1 | from pprint import pprint
+2 2 | 
+3   |-pprint("Hello, world!")  # T203
+4 3 | 
+5 4 | import pprint
+6 5 | 
+
+T203.py:7:1: T203 [*] `pprint` found
   |
 5 | import pprint
 6 | 
@@ -20,5 +29,14 @@ T203.py:7:1: T203 `pprint` found
 8 | 
 9 | pprint.pformat("Hello, world!")
   |
+  = help: Remove `pprint`
+
+ℹ Unsafe fix
+4 4 | 
+5 5 | import pprint
+6 6 | 
+7   |-pprint.pprint("Hello, world!")  # T203
+8 7 | 
+9 8 | pprint.pformat("Hello, world!")
 
 
