From 1bca447fddc179a0d1ea0b4b9c280a2e81e6713e Mon Sep 17 00:00:00 2001 From: Aleksei Latyshev Date: Wed, 27 Mar 2024 17:57:36 +0100 Subject: [PATCH] [`refurb`] Implement `for-loop-writes` (FURB122) lint --- .../resources/test/fixtures/refurb/FURB122.py | 60 ++++++++ .../src/checkers/ast/analyze/statement.rs | 3 + crates/ruff_linter/src/codes.rs | 1 + crates/ruff_linter/src/rules/refurb/mod.rs | 1 + .../src/rules/refurb/rules/for_loop_writes.rs | 121 +++++++++++++++ .../ruff_linter/src/rules/refurb/rules/mod.rs | 2 + ...es__refurb__tests__FURB122_FURB122.py.snap | 142 ++++++++++++++++++ ruff.schema.json | 1 + 8 files changed, 331 insertions(+) create mode 100644 crates/ruff_linter/resources/test/fixtures/refurb/FURB122.py create mode 100644 crates/ruff_linter/src/rules/refurb/rules/for_loop_writes.rs create mode 100644 crates/ruff_linter/src/rules/refurb/snapshots/ruff_linter__rules__refurb__tests__FURB122_FURB122.py.snap diff --git a/crates/ruff_linter/resources/test/fixtures/refurb/FURB122.py b/crates/ruff_linter/resources/test/fixtures/refurb/FURB122.py new file mode 100644 index 00000000000000..29cb44c47e66b7 --- /dev/null +++ b/crates/ruff_linter/resources/test/fixtures/refurb/FURB122.py @@ -0,0 +1,60 @@ +from pathlib import Path + +lines = ["line 1", "line 2", "line 3"] + +# Errors + +with open("file", "w") as f: + for line in lines: + f.write(line) + +other_line = "other line" +with Path("file").open("w") as f: + for line in lines: + f.write(other_line) + +with Path("file").open("w") as f: + for line in lines: + f.write(line) + +with Path("file").open("wb") as f: + for line in lines: + f.write(line.encode()) + +with Path("file").open("w") as f: + for line in lines: + f.write(line.upper()) + +with Path("file").open("w") as f: + pass + + for line in lines: + f.write(line) + +# OK + +for line in lines: + Path("file").open("w").write(line) + +with Path("file").open("w") as f: + for line in lines: + pass + + f.write(line) + +with Path("file").open("w") as f: + for line in lines: + f.write(line) + else: + pass + + +async def func(): + with Path("file").open("w") as f: + async for line in lines: # type: ignore + f.write(line) + + +with Path("file").open("w") as f: + for line in lines: + f.write() # type: ignore diff --git a/crates/ruff_linter/src/checkers/ast/analyze/statement.rs b/crates/ruff_linter/src/checkers/ast/analyze/statement.rs index 0b7839767356d8..6702629505e4a6 100644 --- a/crates/ruff_linter/src/checkers/ast/analyze/statement.rs +++ b/crates/ruff_linter/src/checkers/ast/analyze/statement.rs @@ -1318,6 +1318,9 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) { if checker.enabled(Rule::ForLoopSetMutations) { refurb::rules::for_loop_set_mutations(checker, for_stmt); } + if checker.enabled(Rule::ForLoopWrites) { + refurb::rules::for_loop_writes(checker, for_stmt); + } } } Stmt::Try(ast::StmtTry { diff --git a/crates/ruff_linter/src/codes.rs b/crates/ruff_linter/src/codes.rs index 7b5dbd8a24502e..0fabef9a1c93b2 100644 --- a/crates/ruff_linter/src/codes.rs +++ b/crates/ruff_linter/src/codes.rs @@ -1037,6 +1037,7 @@ pub fn code_to_rule(linter: Linter, code: &str) -> Option<(RuleGroup, Rule)> { #[allow(deprecated)] (Refurb, "113") => (RuleGroup::Nursery, rules::refurb::rules::RepeatedAppend), (Refurb, "118") => (RuleGroup::Preview, rules::refurb::rules::ReimplementedOperator), + (Refurb, "122") => (RuleGroup::Preview, rules::refurb::rules::ForLoopWrites), (Refurb, "129") => (RuleGroup::Preview, rules::refurb::rules::ReadlinesInFor), #[allow(deprecated)] (Refurb, "131") => (RuleGroup::Nursery, rules::refurb::rules::DeleteFullSlice), diff --git a/crates/ruff_linter/src/rules/refurb/mod.rs b/crates/ruff_linter/src/rules/refurb/mod.rs index 1408109b2c6e63..6340fdccf0bec3 100644 --- a/crates/ruff_linter/src/rules/refurb/mod.rs +++ b/crates/ruff_linter/src/rules/refurb/mod.rs @@ -17,6 +17,7 @@ mod tests { #[test_case(Rule::ReadWholeFile, Path::new("FURB101.py"))] #[test_case(Rule::RepeatedAppend, Path::new("FURB113.py"))] #[test_case(Rule::ReimplementedOperator, Path::new("FURB118.py"))] + #[test_case(Rule::ForLoopWrites, Path::new("FURB122.py"))] #[test_case(Rule::ReadlinesInFor, Path::new("FURB129.py"))] #[test_case(Rule::DeleteFullSlice, Path::new("FURB131.py"))] #[test_case(Rule::CheckAndRemoveFromSet, Path::new("FURB132.py"))] diff --git a/crates/ruff_linter/src/rules/refurb/rules/for_loop_writes.rs b/crates/ruff_linter/src/rules/refurb/rules/for_loop_writes.rs new file mode 100644 index 00000000000000..b97885ade95c67 --- /dev/null +++ b/crates/ruff_linter/src/rules/refurb/rules/for_loop_writes.rs @@ -0,0 +1,121 @@ +use ruff_diagnostics::{AlwaysFixableViolation, Diagnostic, Edit, Fix}; +use ruff_macros::{derive_message_formats, violation}; +use ruff_python_ast::{Expr, Stmt, StmtFor}; +use ruff_python_semantic::analyze::typing; + +use crate::checkers::ast::Checker; + +/// ## What it does +/// Checks use of `IOBase.write` in a for loop. +/// +/// ## Why is this bad? +/// When writing a batch of elements, it's more idiomatic to use a single method call +/// `IOBase.writelines` rather than write elements one by one. +/// +/// ## Example +/// ```python +/// with Path("file").open("w") as f: +/// for line in lines: +/// f.write(line) +/// +/// with Path("file").open("wb") as f: +/// for line in lines: +/// f.write(line.encode()) +/// ``` +/// +/// Use instead: +/// ```python +/// with Path("file").open("w") as f: +/// f.writelines(lines) +/// +/// with Path("file").open("wb") as f: +/// f.writelines(line.encode() for line in lines) +/// ``` +/// +/// ## References +/// - [Python documentation: `io.IOBase.writelines`](https://docs.python.org/3/library/io.html#io.IOBase.writelines) +#[violation] +pub struct ForLoopWrites { + name: String, +} + +impl AlwaysFixableViolation for ForLoopWrites { + #[derive_message_formats] + fn message(&self) -> String { + format!("Use of `{}.write` in a for loop", self.name) + } + + fn fix_title(&self) -> String { + format!("Replace with `{}.writelines`", self.name) + } +} + +pub(crate) fn for_loop_writes(checker: &mut Checker, for_stmt: &StmtFor) { + if !for_stmt.orelse.is_empty() { + return; + } + let [Stmt::Expr(stmt_expr)] = for_stmt.body.as_slice() else { + return; + }; + let Expr::Call(call_expr) = stmt_expr.value.as_ref() else { + return; + }; + let Expr::Attribute(expr_attr) = call_expr.func.as_ref() else { + return; + }; + if expr_attr.attr.as_str() != "write" { + return; + } + if !call_expr.arguments.keywords.is_empty() { + return; + } + let [write_arg] = call_expr.arguments.args.as_ref() else { + return; + }; + + let Expr::Name(io_object_name) = expr_attr.value.as_ref() else { + return; + }; + + // Determine whether `f` in `f.write()` was bound to a file object. + if !checker + .semantic() + .resolve_name(io_object_name) + .map(|id| checker.semantic().binding(id)) + .is_some_and(|binding| typing::is_io_base(binding, checker.semantic())) + { + return; + } + + let content = match (for_stmt.target.as_ref(), write_arg) { + (Expr::Name(for_target), Expr::Name(write_arg)) if for_target.id == write_arg.id => { + format!( + "{}.writelines({})", + checker.locator().slice(io_object_name), + checker.locator().slice(for_stmt.iter.as_ref()), + ) + } + (for_target, write_arg) => { + format!( + "{}.writelines({} for {} in {})", + checker.locator().slice(io_object_name), + checker.locator().slice(write_arg), + checker.locator().slice(for_target), + checker.locator().slice(for_stmt.iter.as_ref()), + ) + } + }; + + checker.diagnostics.push( + Diagnostic::new( + ForLoopWrites { + name: io_object_name.id.clone(), + }, + for_stmt.range, + ) + .with_fix(Fix::safe_edit(Edit::range_replacement( + content, + for_stmt.range, + ))), + ); +} diff --git a/crates/ruff_linter/src/rules/refurb/rules/mod.rs b/crates/ruff_linter/src/rules/refurb/rules/mod.rs index 832342c509e13b..9725bb32369998 100644 --- a/crates/ruff_linter/src/rules/refurb/rules/mod.rs +++ b/crates/ruff_linter/src/rules/refurb/rules/mod.rs @@ -2,6 +2,7 @@ pub(crate) use bit_count::*; pub(crate) use check_and_remove_from_set::*; pub(crate) use delete_full_slice::*; pub(crate) use for_loop_set_mutations::*; +pub(crate) use for_loop_writes::*; pub(crate) use hashlib_digest_hex::*; pub(crate) use if_expr_min_max::*; pub(crate) use implicit_cwd::*; @@ -27,6 +28,7 @@ mod bit_count; mod check_and_remove_from_set; mod delete_full_slice; mod for_loop_set_mutations; +mod for_loop_writes; mod hashlib_digest_hex; mod if_expr_min_max; mod implicit_cwd; diff --git a/crates/ruff_linter/src/rules/refurb/snapshots/ruff_linter__rules__refurb__tests__FURB122_FURB122.py.snap b/crates/ruff_linter/src/rules/refurb/snapshots/ruff_linter__rules__refurb__tests__FURB122_FURB122.py.snap new file mode 100644 index 00000000000000..6f3f8e3932dfe6 --- /dev/null +++ b/crates/ruff_linter/src/rules/refurb/snapshots/ruff_linter__rules__refurb__tests__FURB122_FURB122.py.snap @@ -0,0 +1,142 @@ +--- +source: crates/ruff_linter/src/rules/refurb/mod.rs +--- +FURB122.py:8:5: FURB122 [*] Use of `f.write` in a for loop + | + 7 | with open("file", "w") as f: + 8 | for line in lines: + | _____^ + 9 | | f.write(line) + | |_____________________^ FURB122 +10 | +11 | other_line = "other line" + | + = help: Replace with `f.writelines` + +ℹ Safe fix +5 5 | # Errors +6 6 | +7 7 | with open("file", "w") as f: +8 |- for line in lines: +9 |- f.write(line) + 8 |+ f.writelines(lines) +10 9 | +11 10 | other_line = "other line" +12 11 | with Path("file").open("w") as f: + +FURB122.py:13:5: FURB122 [*] Use of `f.write` in a for loop + | +11 | other_line = "other line" +12 | with Path("file").open("w") as f: +13 | for line in lines: + | _____^ +14 | | f.write(other_line) + | |___________________________^ FURB122 +15 | +16 | with Path("file").open("w") as f: + | + = help: Replace with `f.writelines` + +ℹ Safe fix +10 10 | +11 11 | other_line = "other line" +12 12 | with Path("file").open("w") as f: +13 |- for line in lines: +14 |- f.write(other_line) + 13 |+ f.writelines(other_line for line in lines) +15 14 | +16 15 | with Path("file").open("w") as f: +17 16 | for line in lines: + +FURB122.py:17:5: FURB122 [*] Use of `f.write` in a for loop + | +16 | with Path("file").open("w") as f: +17 | for line in lines: + | _____^ +18 | | f.write(line) + | |_____________________^ FURB122 +19 | +20 | with Path("file").open("wb") as f: + | + = help: Replace with `f.writelines` + +ℹ Safe fix +14 14 | f.write(other_line) +15 15 | +16 16 | with Path("file").open("w") as f: +17 |- for line in lines: +18 |- f.write(line) + 17 |+ f.writelines(lines) +19 18 | +20 19 | with Path("file").open("wb") as f: +21 20 | for line in lines: + +FURB122.py:21:5: FURB122 [*] Use of `f.write` in a for loop + | +20 | with Path("file").open("wb") as f: +21 | for line in lines: + | _____^ +22 | | f.write(line.encode()) + | |______________________________^ FURB122 +23 | +24 | with Path("file").open("w") as f: + | + = help: Replace with `f.writelines` + +ℹ Safe fix +18 18 | f.write(line) +19 19 | +20 20 | with Path("file").open("wb") as f: +21 |- for line in lines: +22 |- f.write(line.encode()) + 21 |+ f.writelines(line.encode() for line in lines) +23 22 | +24 23 | with Path("file").open("w") as f: +25 24 | for line in lines: + +FURB122.py:25:5: FURB122 [*] Use of `f.write` in a for loop + | +24 | with Path("file").open("w") as f: +25 | for line in lines: + | _____^ +26 | | f.write(line.upper()) + | |_____________________________^ FURB122 +27 | +28 | with Path("file").open("w") as f: + | + = help: Replace with `f.writelines` + +ℹ Safe fix +22 22 | f.write(line.encode()) +23 23 | +24 24 | with Path("file").open("w") as f: +25 |- for line in lines: +26 |- f.write(line.upper()) + 25 |+ f.writelines(line.upper() for line in lines) +27 26 | +28 27 | with Path("file").open("w") as f: +29 28 | pass + +FURB122.py:31:5: FURB122 [*] Use of `f.write` in a for loop + | +29 | pass +30 | +31 | for line in lines: + | _____^ +32 | | f.write(line) + | |_____________________^ FURB122 +33 | +34 | # OK + | + = help: Replace with `f.writelines` + +ℹ Safe fix +28 28 | with Path("file").open("w") as f: +29 29 | pass +30 30 | +31 |- for line in lines: +32 |- f.write(line) + 31 |+ f.writelines(lines) +33 32 | +34 33 | # OK +35 34 | diff --git a/ruff.schema.json b/ruff.schema.json index a0982b2c876908..b64e96725cdf9b 100644 --- a/ruff.schema.json +++ b/ruff.schema.json @@ -3007,6 +3007,7 @@ "FURB113", "FURB118", "FURB12", + "FURB122", "FURB129", "FURB13", "FURB131",