Allow f-string modifications in line-shrinking cases (#7818)

## Summary

This PR resolves an issue raised in
#7810, whereby we don't fix
an f-string that exceeds the line length _even if_ the resultant code is
_shorter_ than the current code.

As part of this change, I've also refactored and extracted some common
logic we use around "ensuring a fix isn't breaking the line length
rules".

## Test Plan

`cargo test`

crates/ruff_linter/resources/test/fixtures/pyupgrade/UP032_0.py

@@ -202,3 +202,8 @@ async def c():
 "{}".format(
     1  # comment
 )
+
+
+# The fixed string will exceed the line length, but it's still smaller than the
+# existing line length, so it's fine.
+"<Customer: {}, {}, {}, {}, {}>".format(self.internal_ids, self.external_ids, self.properties, self.tags, self.others)

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

@@ -414,13 +414,7 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
                                         pyupgrade::rules::format_literals(checker, call, &summary);
                                     }
                                     if checker.enabled(Rule::FString) {
-                                        pyupgrade::rules::f_strings(
-                                            checker,
-                                            call,
-                                            &summary,
-                                            value,
-                                            checker.settings.line_length,
-                                        );
+                                        pyupgrade::rules::f_strings(checker, call, &summary, value);
                                     }
                                 }
                             }

crates/ruff_linter/src/fix/edits.rs

@@ -3,16 +3,18 @@
 use anyhow::{Context, Result};
 
 use ruff_diagnostics::Edit;
+use ruff_python_ast::node::AnyNodeRef;
 use ruff_python_ast::{self as ast, Arguments, ExceptHandler, Stmt};
 use ruff_python_codegen::Stylist;
 use ruff_python_index::Indexer;
 use ruff_python_trivia::{
     has_leading_content, is_python_whitespace, PythonWhitespace, SimpleTokenKind, SimpleTokenizer,
 };
-use ruff_source_file::{Locator, NewlineWithTrailingNewline};
+use ruff_source_file::{Locator, NewlineWithTrailingNewline, UniversalNewlines};
 use ruff_text_size::{Ranged, TextLen, TextRange, TextSize};
 
 use crate::fix::codemods;
+use crate::line_width::{LineLength, LineWidthBuilder, TabSize};
 
 /// Return the `Fix` to use when deleting a `Stmt`.
 ///
@@ -285,6 +287,75 @@ pub(crate) fn pad(content: String, range: TextRange, locator: &Locator) -> Strin
     )
 }
 
+/// Returns `true` if the fix fits within the maximum configured line length.
+pub(crate) fn fits(
+    fix: &str,
+    node: AnyNodeRef,
+    locator: &Locator,
+    line_length: LineLength,
+    tab_size: TabSize,
+) -> bool {
+    all_lines_fit(fix, node, locator, line_length.value() as usize, tab_size)
+}
+
+/// Returns `true` if the fix fits within the maximum configured line length, or produces lines that
+/// are shorter than the maximum length of the existing AST node.
+pub(crate) fn fits_or_shrinks(
+    fix: &str,
+    node: AnyNodeRef,
+    locator: &Locator,
+    line_length: LineLength,
+    tab_size: TabSize,
+) -> bool {
+    // Use the larger of the line length limit, or the longest line in the existing AST node.
+    let line_length = std::iter::once(line_length.value() as usize)
+        .chain(
+            locator
+                .slice(locator.lines_range(node.range()))
+                .universal_newlines()
+                .map(|line| LineWidthBuilder::new(tab_size).add_str(&line).get()),
+        )
+        .max()
+        .unwrap_or(line_length.value() as usize);
+
+    all_lines_fit(fix, node, locator, line_length, tab_size)
+}
+
+/// Returns `true` if all lines in the fix are shorter than the given line length.
+fn all_lines_fit(
+    fix: &str,
+    node: AnyNodeRef,
+    locator: &Locator,
+    line_length: usize,
+    tab_size: TabSize,
+) -> bool {
+    let prefix = locator.slice(TextRange::new(
+        locator.line_start(node.start()),
+        node.start(),
+    ));
+
+    // Ensure that all lines are shorter than the line length limit.
+    fix.universal_newlines().enumerate().all(|(idx, line)| {
+        // If `template` is a multiline string, `col_offset` should only be applied to the first
+        // line:
+        // ```
+        // a = """{}        -> offset = col_offset (= 4)
+        // {}               -> offset = 0
+        // """.format(0, 1) -> offset = 0
+        // ```
+        let measured_length = if idx == 0 {
+            LineWidthBuilder::new(tab_size)
+                .add_str(prefix)
+                .add_str(&line)
+                .get()
+        } else {
+            LineWidthBuilder::new(tab_size).add_str(&line).get()
+        };
+
+        measured_length <= line_length
+    })
+}
+
 #[cfg(test)]
 mod tests {
     use anyhow::Result;

crates/ruff_linter/src/line_width.rs

@@ -1,12 +1,14 @@
-use ruff_cache::{CacheKey, CacheKeyHasher};
-use serde::{Deserialize, Serialize};
 use std::error::Error;
 use std::hash::Hasher;
 use std::num::{NonZeroU16, NonZeroU8, ParseIntError};
 use std::str::FromStr;
+
+use serde::{Deserialize, Serialize};
 use unicode_width::UnicodeWidthChar;
 
+use ruff_cache::{CacheKey, CacheKeyHasher};
 use ruff_macros::CacheKey;
+use ruff_text_size::TextSize;
 
 /// The length of a line of text that is considered too long.
 ///
@@ -23,6 +25,10 @@ impl LineLength {
     pub fn value(&self) -> u16 {
         self.0.get()
     }
+
+    pub fn text_len(&self) -> TextSize {
+        TextSize::from(u32::from(self.value()))
+    }
 }
 
 impl Default for LineLength {

crates/ruff_linter/src/rules/flake8_simplify/rules/ast_if.rs

@@ -12,11 +12,11 @@ use ruff_python_ast::{
 };
 use ruff_python_semantic::SemanticModel;
 use ruff_python_trivia::{SimpleTokenKind, SimpleTokenizer};
-use ruff_source_file::{Locator, UniversalNewlines};
+use ruff_source_file::Locator;
 use ruff_text_size::{Ranged, TextRange};
 
 use crate::checkers::ast::Checker;
-use crate::line_width::LineWidthBuilder;
+use crate::fix::edits::fits;
 use crate::registry::AsRule;
 use crate::rules::flake8_simplify::rules::fix_if;
 
@@ -443,15 +443,15 @@ pub(crate) fn nested_if_statements(
             match fix_if::fix_nested_if_statements(checker.locator(), checker.stylist(), nested_if)
             {
                 Ok(edit) => {
-                    if edit
-                        .content()
-                        .unwrap_or_default()
-                        .universal_newlines()
-                        .all(|line| {
-                            LineWidthBuilder::new(checker.settings.tab_size).add_str(&line)
-                                <= checker.settings.line_length
-                        })
-                    {
+                    if edit.content().map_or(true, |content| {
+                        fits(
+                            content,
+                            (&nested_if).into(),
+                            checker.locator(),
+                            checker.settings.line_length,
+                            checker.settings.tab_size,
+                        )
+                    }) {
                         diagnostic.set_fix(Fix::suggested(edit));
                     }
                 }
@@ -693,16 +693,13 @@ pub(crate) fn use_ternary_operator(checker: &mut Checker, stmt: &Stmt) {
     let contents = checker.generator().stmt(&ternary);
 
     // Don't flag if the resulting expression would exceed the maximum line length.
-    let line_start = checker.locator().line_start(stmt.start());
-    if LineWidthBuilder::new(checker.settings.tab_size)
-        .add_str(
-            checker
-                .locator()
-                .slice(TextRange::new(line_start, stmt.start())),
-        )
-        .add_str(&contents)
-        > checker.settings.line_length
-    {
+    if !fits(
+        &contents,
+        stmt.into(),
+        checker.locator(),
+        checker.settings.line_length,
+        checker.settings.tab_size,
+    ) {
         return;
     }
 
@@ -1001,16 +998,13 @@ pub(crate) fn use_dict_get_with_default(checker: &mut Checker, stmt_if: &ast::St
     let contents = checker.generator().stmt(&node5.into());
 
     // Don't flag if the resulting expression would exceed the maximum line length.
-    let line_start = checker.locator().line_start(stmt_if.start());
-    if LineWidthBuilder::new(checker.settings.tab_size)
-        .add_str(
-            checker
-                .locator()
-                .slice(TextRange::new(line_start, stmt_if.start())),
-        )
-        .add_str(&contents)
-        > checker.settings.line_length
-    {
+    if !fits(
+        &contents,
+        stmt_if.into(),
+        checker.locator(),
+        checker.settings.line_length,
+        checker.settings.tab_size,
+    ) {
         return;
     }
 

crates/ruff_linter/src/rules/flake8_simplify/rules/ast_with.rs

@@ -5,11 +5,10 @@ use ruff_diagnostics::{FixAvailability, Violation};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::{self as ast, Stmt, WithItem};
 use ruff_python_trivia::{SimpleTokenKind, SimpleTokenizer};
-use ruff_source_file::UniversalNewlines;
 use ruff_text_size::{Ranged, TextRange};
 
 use crate::checkers::ast::Checker;
-use crate::line_width::LineWidthBuilder;
+use crate::fix::edits::fits;
 use crate::registry::AsRule;
 
 use super::fix_with;
@@ -137,15 +136,15 @@ pub(crate) fn multiple_with_statements(
                     with_stmt,
                 ) {
                     Ok(edit) => {
-                        if edit
-                            .content()
-                            .unwrap_or_default()
-                            .universal_newlines()
-                            .all(|line| {
-                                LineWidthBuilder::new(checker.settings.tab_size).add_str(&line)
-                                    <= checker.settings.line_length
-                            })
-                        {
+                        if edit.content().map_or(true, |content| {
+                            fits(
+                                content,
+                                with_stmt.into(),
+                                checker.locator(),
+                                checker.settings.line_length,
+                                checker.settings.tab_size,
+                            )
+                        }) {
                             diagnostic.set_fix(Fix::suggested(edit));
                         }
                     }

crates/ruff_linter/src/rules/flake8_simplify/rules/reimplemented_builtin.rs

@@ -1,15 +1,15 @@
-use ruff_python_ast::{
-    self as ast, Arguments, CmpOp, Comprehension, Constant, Expr, ExprContext, Stmt, UnaryOp,
-};
-use ruff_text_size::{Ranged, TextRange};
-
 use ruff_diagnostics::{Diagnostic, Edit, Fix, FixAvailability, Violation};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast::helpers::any_over_expr;
 use ruff_python_ast::traversal;
+use ruff_python_ast::{
+    self as ast, Arguments, CmpOp, Comprehension, Constant, Expr, ExprContext, Stmt, UnaryOp,
+};
 use ruff_python_codegen::Generator;
+use ruff_text_size::{Ranged, TextRange};
 
 use crate::checkers::ast::Checker;
+use crate::fix::edits::fits;
 use crate::line_width::LineWidthBuilder;
 use crate::registry::AsRule;
 
@@ -98,16 +98,13 @@ pub(crate) fn convert_for_loop_to_any_all(checker: &mut Checker, stmt: &Stmt) {
             );
 
             // Don't flag if the resulting expression would exceed the maximum line length.
-            let line_start = checker.locator().line_start(stmt.start());
-            if LineWidthBuilder::new(checker.settings.tab_size)
-                .add_str(
-                    checker
-                        .locator()
-                        .slice(TextRange::new(line_start, stmt.start())),
-                )
-                .add_str(&contents)
-                > checker.settings.line_length
-            {
+            if !fits(
+                &contents,
+                stmt.into(),
+                checker.locator(),
+                checker.settings.line_length,
+                checker.settings.tab_size,
+            ) {
                 return;
             }
 

crates/ruff_linter/src/rules/pyupgrade/rules/f_strings.rs

@@ -11,12 +11,12 @@ use ruff_python_literal::format::{
     FieldName, FieldNamePart, FieldType, FormatPart, FormatString, FromTemplate,
 };
 use ruff_python_parser::{lexer, Mode, Tok};
-
 use ruff_source_file::Locator;
 use ruff_text_size::{Ranged, TextRange};
 
 use crate::checkers::ast::Checker;
-use crate::line_width::LineLength;
+use crate::fix::edits::fits_or_shrinks;
+
 use crate::registry::AsRule;
 use crate::rules::pyflakes::format::FormatSummary;
 use crate::rules::pyupgrade::helpers::curly_escape;
@@ -306,7 +306,6 @@ pub(crate) fn f_strings(
     call: &ast::ExprCall,
     summary: &FormatSummary,
     template: &Expr,
-    line_length: LineLength,
 ) {
     if summary.has_nested_parts {
         return;
@@ -384,22 +383,6 @@ pub(crate) fn f_strings(
     }
     contents.push_str(checker.locator().slice(TextRange::new(prev_end, end)));
 
-    // Avoid refactors that exceed the line length limit.
-    let col_offset = template.start() - checker.locator().line_start(template.start());
-    if contents.lines().enumerate().any(|(idx, line)| {
-        // If `template` is a multiline string, `col_offset` should only be applied to the first
-        // line:
-        // ```
-        // a = """{}        -> offset = col_offset (= 4)
-        // {}               -> offset = 0
-        // """.format(0, 1) -> offset = 0
-        // ```
-        let offset = if idx == 0 { col_offset.to_usize() } else { 0 };
-        offset + line.chars().count() > line_length.value() as usize
-    }) {
-        return;
-    }
-
     // If necessary, add a space between any leading keyword (`return`, `yield`, `assert`, etc.)
     // and the string. For example, `return"foo"` is valid, but `returnf"foo"` is not.
     let existing = checker.locator().slice(TextRange::up_to(call.start()));
@@ -411,6 +394,17 @@ pub(crate) fn f_strings(
         contents.insert(0, ' ');
     }
 
+    // Avoid refactors that exceed the line length limit.
+    if !fits_or_shrinks(
+        &contents,
+        template.into(),
+        checker.locator(),
+        checker.settings.line_length,
+        checker.settings.tab_size,
+    ) {
+        return;
+    }
+
     let mut diagnostic = Diagnostic::new(FString, call.range());
 
     // Avoid fix if there are comments within the call: