[refurb] Implement for-loop-writes (FURB122)

crates/ruff_linter/resources/test/fixtures/refurb/FURB122.py

@@ -0,0 +1,66 @@
+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)
+
+# Offer unsafe fix if it would delete comments
+with open("file","w") as f:
+    for line in lines:
+        # a really important comment
+        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

crates/ruff_linter/src/checkers/ast/analyze/statement.rs

@@ -1431,6 +1431,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 {

crates/ruff_linter/src/codes.rs

@@ -1089,6 +1089,7 @@ pub fn code_to_rule(linter: Linter, code: &str) -> Option<(RuleGroup, Rule)> {
         (Refurb, "113") => (RuleGroup::Preview, rules::refurb::rules::RepeatedAppend),
         (Refurb, "116") => (RuleGroup::Preview, rules::refurb::rules::FStringNumberFormat),
         (Refurb, "118") => (RuleGroup::Preview, rules::refurb::rules::ReimplementedOperator),
+        (Refurb, "122") => (RuleGroup::Preview, rules::refurb::rules::ForLoopWrites),
         (Refurb, "129") => (RuleGroup::Stable, rules::refurb::rules::ReadlinesInFor),
         (Refurb, "131") => (RuleGroup::Preview, rules::refurb::rules::DeleteFullSlice),
         (Refurb, "132") => (RuleGroup::Preview, rules::refurb::rules::CheckAndRemoveFromSet),

crates/ruff_linter/src/rules/refurb/mod.rs

@@ -19,6 +19,7 @@ mod tests {
     #[test_case(Rule::RepeatedAppend, Path::new("FURB113.py"))]
     #[test_case(Rule::IfExpInsteadOfOrOperator, Path::new("FURB110.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"))]

crates/ruff_linter/src/rules/refurb/rules/for_loop_writes.rs

@@ -0,0 +1,128 @@
+use ruff_diagnostics::{AlwaysFixableViolation, Applicability, Diagnostic, Edit, Fix};
+use ruff_macros::{derive_message_formats, ViolationMetadata};
+use ruff_python_ast::{Expr, Stmt, StmtFor};
+use ruff_python_semantic::analyze::typing;
+use ruff_text_size::Ranged;
+
+use crate::checkers::ast::Checker;
+
+/// ## What it does
+/// Checks for the 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 to
+/// `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)
+#[derive(ViolationMetadata)]
+pub(crate) 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()),
+            )
+        }
+    };
+
+    let applicability = if checker.comment_ranges().intersects(for_stmt.range()) {
+        Applicability::Unsafe
+    } else {
+        Applicability::Safe
+    };
+
+    checker.diagnostics.push(
+        Diagnostic::new(
+            ForLoopWrites {
+                name: io_object_name.id.to_string(),
+            },
+            for_stmt.range,
+        )
+        .with_fix(Fix::applicable_edit(
+            Edit::range_replacement(content, for_stmt.range),
+            applicability,
+        )),
+    );
+}

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 fstring_number_format::*;
 pub(crate) use hardcoded_string_charset::*;
 pub(crate) use hashlib_digest_hex::*;
@@ -37,6 +38,7 @@ mod bit_count;
 mod check_and_remove_from_set;
 mod delete_full_slice;
 mod for_loop_set_mutations;
+mod for_loop_writes;
 mod fstring_number_format;
 mod hardcoded_string_charset;
 mod hashlib_digest_hex;

crates/ruff_linter/src/rules/refurb/snapshots/ruff_linter__rules__refurb__tests__FURB122_FURB122.py.snap

@@ -0,0 +1,161 @@
+---
+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 |   # Offer unsafe fix if it would delete comments
+   |
+   = 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 | # Offer unsafe fix if it would delete comments
+35 34 | with open("file","w") as f:
+
+FURB122.py:36:5: FURB122 [*] Use of `f.write` in a for loop
+   |
+34 |   # Offer unsafe fix if it would delete comments
+35 |   with open("file","w") as f:
+36 | /     for line in lines:
+37 | |         # a really important comment
+38 | |         f.write(line)
+   | |_____________________^ FURB122
+39 |
+40 |   # OK
+   |
+   = help: Replace with `f.writelines`
+
+ℹ Unsafe fix
+33 33 | 
+34 34 | # Offer unsafe fix if it would delete comments
+35 35 | with open("file","w") as f:
+36    |-    for line in lines:
+37    |-        # a really important comment
+38    |-        f.write(line)
+   36 |+    f.writelines(lines)
+39 37 | 
+40 38 | # OK
+41 39 |