Format implicit string continuation (#5328)

crates/ruff_python_formatter/src/builders.rs

@@ -1,10 +1,35 @@
 use crate::context::NodeLevel;
 use crate::prelude::*;
 use crate::trivia::{first_non_trivia_token, lines_after, skip_trailing_trivia, Token, TokenKind};
-use ruff_formatter::write;
+use ruff_formatter::{format_args, write, Argument, Arguments};
 use ruff_text_size::TextSize;
 use rustpython_parser::ast::Ranged;
 
+/// Adds parentheses and indents `content` if it doesn't fit on a line.
+pub(crate) fn optional_parentheses<'ast, T>(content: &T) -> OptionalParentheses<'_, 'ast>
+where
+    T: Format<PyFormatContext<'ast>>,
+{
+    OptionalParentheses {
+        inner: Argument::new(content),
+    }
+}
+
+pub(crate) struct OptionalParentheses<'a, 'ast> {
+    inner: Argument<'a, PyFormatContext<'ast>>,
+}
+
+impl<'ast> Format<PyFormatContext<'ast>> for OptionalParentheses<'_, 'ast> {
+    fn fmt(&self, f: &mut Formatter<PyFormatContext<'ast>>) -> FormatResult<()> {
+        group(&format_args![
+            if_group_breaks(&text("(")),
+            soft_block_indent(&Arguments::from(&self.inner)),
+            if_group_breaks(&text(")"))
+        ])
+        .fmt(f)
+    }
+}
+
 /// Provides Python specific extensions to [`Formatter`].
 pub(crate) trait PyFormatterExtensions<'ast, 'buf> {
     /// Creates a joiner that inserts the appropriate number of empty lines between two nodes, depending on the

crates/ruff_python_formatter/src/expression/expr_constant.rs

@@ -2,14 +2,25 @@ use crate::comments::Comments;
 use crate::expression::parentheses::{
     default_expression_needs_parentheses, NeedsParentheses, Parentheses, Parenthesize,
 };
-use crate::expression::string::FormatString;
+use crate::expression::string::{FormatString, StringLayout};
 use crate::prelude::*;
 use crate::{not_yet_implemented_custom_text, verbatim_text, FormatNodeRule};
-use ruff_formatter::write;
+use ruff_formatter::{write, FormatRuleWithOptions};
 use rustpython_parser::ast::{Constant, ExprConstant};
 
 #[derive(Default)]
-pub struct FormatExprConstant;
+pub struct FormatExprConstant {
+    string_layout: StringLayout,
+}
+
+impl FormatRuleWithOptions<ExprConstant, PyFormatContext<'_>> for FormatExprConstant {
+    type Options = StringLayout;
+
+    fn with_options(mut self, options: Self::Options) -> Self {
+        self.string_layout = options;
+        self
+    }
+}
 
 impl FormatNodeRule<ExprConstant> for FormatExprConstant {
     fn fmt_fields(&self, item: &ExprConstant, f: &mut PyFormatter) -> FormatResult<()> {
@@ -29,7 +40,7 @@ impl FormatNodeRule<ExprConstant> for FormatExprConstant {
             Constant::Int(_) | Constant::Float(_) | Constant::Complex { .. } => {
                 write!(f, [verbatim_text(item)])
             }
-            Constant::Str(_) => FormatString::new(item).fmt(f),
+            Constant::Str(_) => FormatString::new(item, self.string_layout).fmt(f),
             Constant::Bytes(_) => {
                 not_yet_implemented_custom_text(r#"b"NOT_YET_IMPLEMENTED_BYTE_STRING""#).fmt(f)
             }
@@ -44,14 +55,6 @@ impl FormatNodeRule<ExprConstant> for FormatExprConstant {
         _node: &ExprConstant,
         _f: &mut PyFormatter,
     ) -> FormatResult<()> {
-        // TODO(konstin): Reactivate when string formatting works, currently a source of unstable
-        // formatting, e.g.:
-        // magic_methods = (
-        //     "enter exit "
-        //     # we added divmod and rdivmod here instead of numerics
-        //     # because there is no idivmod
-        //     "divmod rdivmod neg pos abs invert "
-        // )
         Ok(())
     }
 }
@@ -64,6 +67,14 @@ impl NeedsParentheses for ExprConstant {
         comments: &Comments,
     ) -> Parentheses {
         match default_expression_needs_parentheses(self.into(), parenthesize, source, comments) {
+            Parentheses::Optional if self.value.is_str() && parenthesize.is_if_breaks() => {
+                // Custom handling that only adds parentheses for implicit concatenated strings.
+                if parenthesize.is_if_breaks() {
+                    Parentheses::Custom
+                } else {
+                    Parentheses::Optional
+                }
+            }
             Parentheses::Optional => Parentheses::Never,
             parentheses => parentheses,
         }

crates/ruff_python_formatter/src/expression/expr_tuple.rs

@@ -1,13 +1,10 @@
-use crate::builders::PyFormatterExtensions;
+use crate::builders::optional_parentheses;
 use crate::comments::{dangling_node_comments, Comments};
-use crate::context::PyFormatContext;
 use crate::expression::parentheses::{
     default_expression_needs_parentheses, NeedsParentheses, Parentheses, Parenthesize,
 };
-use crate::{AsFormat, FormatNodeRule, PyFormatter};
-use ruff_formatter::formatter::Formatter;
-use ruff_formatter::prelude::{block_indent, group, if_group_breaks, soft_block_indent, text};
-use ruff_formatter::{format_args, write, Buffer, Format, FormatResult, FormatRuleWithOptions};
+use crate::prelude::*;
+use ruff_formatter::{format_args, write, FormatRuleWithOptions};
 use ruff_python_ast::prelude::{Expr, Ranged};
 use ruff_text_size::TextRange;
 use rustpython_parser::ast::ExprTuple;
@@ -100,17 +97,7 @@ impl FormatNodeRule<ExprTuple> for FormatExprTuple {
                     ])]
                 )
             }
-            elts => {
-                write!(
-                    f,
-                    [group(&format_args![
-                        // If there were previously no parentheses, add them only if the group breaks
-                        if_group_breaks(&text("(")),
-                        soft_block_indent(&ExprSequence::new(elts)),
-                        if_group_breaks(&text(")")),
-                    ])]
-                )
-            }
+            elts => optional_parentheses(&ExprSequence::new(elts)).fmt(f),
         }
     }
 

crates/ruff_python_formatter/src/expression/mod.rs

@@ -1,7 +1,9 @@
+use crate::builders::optional_parentheses;
 use crate::comments::Comments;
 use crate::context::NodeLevel;
 use crate::expression::expr_tuple::TupleParentheses;
 use crate::expression::parentheses::{NeedsParentheses, Parentheses, Parenthesize};
+use crate::expression::string::StringLayout;
 use crate::prelude::*;
 use ruff_formatter::{
     format_args, write, FormatOwnedWithRule, FormatRefWithRule, FormatRule, FormatRuleWithOptions,
@@ -37,7 +39,7 @@ pub(crate) mod expr_unary_op;
 pub(crate) mod expr_yield;
 pub(crate) mod expr_yield_from;
 pub(crate) mod parentheses;
-mod string;
+pub(crate) mod string;
 
 #[derive(Default)]
 pub struct FormatExpr {
@@ -81,7 +83,10 @@ impl FormatRule<Expr, PyFormatContext<'_>> for FormatExpr {
             Expr::Call(expr) => expr.format().fmt(f),
             Expr::FormattedValue(expr) => expr.format().fmt(f),
             Expr::JoinedStr(expr) => expr.format().fmt(f),
-            Expr::Constant(expr) => expr.format().fmt(f),
+            Expr::Constant(expr) => expr
+                .format()
+                .with_options(StringLayout::Default(Some(parentheses)))
+                .fmt(f),
             Expr::Attribute(expr) => expr.format().fmt(f),
             Expr::Subscript(expr) => expr.format().fmt(f),
             Expr::Starred(expr) => expr.format().fmt(f),
@@ -109,16 +114,7 @@ impl FormatRule<Expr, PyFormatContext<'_>> for FormatExpr {
                 )
             }
             // Add optional parentheses. Ignore if the item renders parentheses itself.
-            Parentheses::Optional => {
-                write!(
-                    f,
-                    [group(&format_args![
-                        if_group_breaks(&text("(")),
-                        soft_block_indent(&format_expr),
-                        if_group_breaks(&text(")"))
-                    ])]
-                )
-            }
+            Parentheses::Optional => optional_parentheses(&format_expr).fmt(f),
             Parentheses::Custom | Parentheses::Never => Format::fmt(&format_expr, f),
         };
 

crates/ruff_python_formatter/src/expression/parentheses.rs

@@ -104,7 +104,7 @@ pub enum Parentheses {
     Never,
 }
 
-fn is_expression_parenthesized(expr: AnyNodeRef, contents: &str) -> bool {
+pub(crate) fn is_expression_parenthesized(expr: AnyNodeRef, contents: &str) -> bool {
     matches!(
         first_non_trivia_token(expr.end(), contents),
         Some(Token {

crates/ruff_python_formatter/src/expression/string.rs

@@ -1,35 +1,152 @@
+use crate::builders::optional_parentheses;
+use crate::comments::{leading_comments, trailing_comments};
+use crate::expression::parentheses::Parentheses;
 use crate::prelude::*;
-use crate::{not_yet_implemented_custom_text, QuoteStyle};
+use crate::QuoteStyle;
 use bitflags::bitflags;
-use ruff_formatter::{write, FormatError};
+use ruff_formatter::{format_args, write, FormatError};
 use ruff_python_ast::str::is_implicit_concatenation;
 use ruff_text_size::{TextLen, TextRange, TextSize};
 use rustpython_parser::ast::{ExprConstant, Ranged};
+use rustpython_parser::lexer::lex_starts_at;
+use rustpython_parser::{Mode, Tok};
 use std::borrow::Cow;
 
-pub(super) struct FormatString {
-    string_range: TextRange,
+#[derive(Copy, Clone, Debug)]
+pub enum StringLayout {
+    Default(Option<Parentheses>),
+
+    /// Enforces that implicit continuation strings are printed on a single line even if they exceed
+    /// the configured line width.  
+    Flat,
 }
 
-impl FormatString {
-    pub(super) fn new(constant: &ExprConstant) -> Self {
+impl Default for StringLayout {
+    fn default() -> Self {
+        Self::Default(None)
+    }
+}
+
+pub(super) struct FormatString<'a> {
+    constant: &'a ExprConstant,
+    layout: StringLayout,
+}
+
+impl<'a> FormatString<'a> {
+    pub(super) fn new(constant: &'a ExprConstant, layout: StringLayout) -> Self {
         debug_assert!(constant.value.is_str());
-        Self {
-            string_range: constant.range(),
+        Self { constant, layout }
+    }
+}
+
+impl<'a> Format<PyFormatContext<'_>> for FormatString<'a> {
+    fn fmt(&self, f: &mut Formatter<PyFormatContext<'_>>) -> FormatResult<()> {
+        let string_range = self.constant.range();
+        let string_content = f.context().locator().slice(string_range);
+
+        if is_implicit_concatenation(string_content) {
+            let format_continuation = FormatStringContinuation::new(self.constant, self.layout);
+
+            if let StringLayout::Default(Some(Parentheses::Custom)) = self.layout {
+                optional_parentheses(&format_continuation).fmt(f)
+            } else {
+                format_continuation.fmt(f)
+            }
+        } else {
+            FormatStringPart::new(string_range).fmt(f)
         }
     }
 }
 
-impl Format<PyFormatContext<'_>> for FormatString {
-    fn fmt(&self, f: &mut Formatter<PyFormatContext<'_>>) -> FormatResult<()> {
-        let string_content = f.context().locator().slice(self.string_range);
+struct FormatStringContinuation<'a> {
+    constant: &'a ExprConstant,
+    layout: StringLayout,
+}
 
-        if is_implicit_concatenation(string_content) {
-            not_yet_implemented_custom_text(r#""NOT_YET_IMPLEMENTED" "IMPLICIT_CONCATENATION""#)
-                .fmt(f)
-        } else {
-            FormatStringPart::new(self.string_range).fmt(f)
+impl<'a> FormatStringContinuation<'a> {
+    fn new(constant: &'a ExprConstant, layout: StringLayout) -> Self {
+        debug_assert!(constant.value.is_str());
+        Self { constant, layout }
+    }
+}
+
+impl Format<PyFormatContext<'_>> for FormatStringContinuation<'_> {
+    fn fmt(&self, f: &mut Formatter<PyFormatContext<'_>>) -> FormatResult<()> {
+        let comments = f.context().comments().clone();
+        let locator = f.context().locator();
+        let mut dangling_comments = comments.dangling_comments(self.constant);
+
+        let string_range = self.constant.range();
+        let string_content = locator.slice(string_range);
+
+        // The AST parses implicit concatenation as a single string.
+        // Call into the lexer to extract the individual chunks and format each string on its own.
+        // This code does not yet implement the automatic joining of strings that fit on the same line
+        // because this is a black preview style.
+        let lexer = lex_starts_at(string_content, Mode::Module, string_range.start());
+
+        let separator = format_with(|f| match self.layout {
+            StringLayout::Default(_) => soft_line_break_or_space().fmt(f),
+            StringLayout::Flat => space().fmt(f),
+        });
+
+        let mut joiner = f.join_with(separator);
+
+        for token in lexer {
+            let (token, token_range) = token.map_err(|_| FormatError::SyntaxError)?;
+
+            match token {
+                Tok::String { .. } => {
+                    // ```python
+                    // (
+                    //      "a"
+                    //      # leading
+                    //      "the comment above"
+                    // )
+                    // ```
+                    let leading_comments_end = dangling_comments
+                        .partition_point(|comment| comment.slice().start() <= token_range.start());
+
+                    let (leading_part_comments, rest) =
+                        dangling_comments.split_at(leading_comments_end);
+
+                    // ```python
+                    // (
+                    //      "a" # trailing comment
+                    //      "the comment above"
+                    // )
+                    // ```
+                    let trailing_comments_end = rest.partition_point(|comment| {
+                        comment.line_position().is_end_of_line()
+                            && !locator.contains_line_break(TextRange::new(
+                                token_range.end(),
+                                comment.slice().start(),
+                            ))
+                    });
+
+                    let (trailing_part_comments, rest) = rest.split_at(trailing_comments_end);
+
+                    joiner.entry(&format_args![
+                        line_suffix_boundary(),
+                        leading_comments(leading_part_comments),
+                        FormatStringPart::new(token_range),
+                        trailing_comments(trailing_part_comments)
+                    ]);
+
+                    dangling_comments = rest;
+                }
+                Tok::Comment(_)
+                | Tok::NonLogicalNewline
+                | Tok::Newline
+                | Tok::Indent
+                | Tok::Dedent => continue,
+                token => unreachable!("Unexpected token {token:?}"),
+            }
         }
+
+        debug_assert!(dangling_comments.is_empty());
+
+        joiner.finish()
     }
 }
 

crates/ruff_python_formatter/src/statement/stmt_expr.rs

@@ -1,4 +1,5 @@
-use crate::expression::parentheses::Parenthesize;
+use crate::expression::parentheses::{is_expression_parenthesized, Parenthesize};
+use crate::expression::string::StringLayout;
 use crate::prelude::*;
 use crate::FormatNodeRule;
 use rustpython_parser::ast::StmtExpr;
@@ -10,6 +11,14 @@ impl FormatNodeRule<StmtExpr> for FormatStmtExpr {
     fn fmt_fields(&self, item: &StmtExpr, f: &mut PyFormatter) -> FormatResult<()> {
         let StmtExpr { value, .. } = item;
 
+        if let Some(constant) = value.as_constant_expr() {
+            if constant.value.is_str()
+                && !is_expression_parenthesized(value.as_ref().into(), f.context().contents())
+            {
+                return constant.format().with_options(StringLayout::Flat).fmt(f);
+            }
+        }
+
         value.format().with_options(Parenthesize::Optional).fmt(f)
     }
 }