Associate a trailing end-of-line comment in a parenthesized implicit concatenated string with the last literal (#15378)

Author: MichaReiser

crates/ruff_python_ast/src/expression.rs

@@ -407,13 +407,13 @@ pub enum StringLike<'a> {
     FString(&'a ast::ExprFString),
 }
 
-impl StringLike<'_> {
+impl<'a> StringLike<'a> {
     pub const fn is_fstring(self) -> bool {
         matches!(self, Self::FString(_))
     }
 
     /// Returns an iterator over the [`StringLikePart`] contained in this string-like expression.
-    pub fn parts(&self) -> StringLikePartIter<'_> {
+    pub fn parts(&self) -> StringLikePartIter<'a> {
         match self {
             StringLike::String(expr) => StringLikePartIter::String(expr.value.iter()),
             StringLike::Bytes(expr) => StringLikePartIter::Bytes(expr.value.iter()),
@@ -429,6 +429,14 @@ impl StringLike<'_> {
             Self::FString(ExprFString { value, .. }) => value.is_implicit_concatenated(),
         }
     }
+
+    pub const fn as_expression_ref(self) -> ExpressionRef<'a> {
+        match self {
+            StringLike::String(expr) => ExpressionRef::StringLiteral(expr),
+            StringLike::Bytes(expr) => ExpressionRef::BytesLiteral(expr),
+            StringLike::FString(expr) => ExpressionRef::FString(expr),
+        }
+    }
 }
 
 impl<'a> From<&'a ast::ExprStringLiteral> for StringLike<'a> {
@@ -488,6 +496,19 @@ impl<'a> TryFrom<&'a Expr> for StringLike<'a> {
     }
 }
 
+impl<'a> TryFrom<AnyNodeRef<'a>> for StringLike<'a> {
+    type Error = ();
+
+    fn try_from(value: AnyNodeRef<'a>) -> Result<Self, Self::Error> {
+        match value {
+            AnyNodeRef::ExprStringLiteral(value) => Ok(Self::String(value)),
+            AnyNodeRef::ExprBytesLiteral(value) => Ok(Self::Bytes(value)),
+            AnyNodeRef::ExprFString(value) => Ok(Self::FString(value)),
+            _ => Err(()),
+        }
+    }
+}
+
 impl Ranged for StringLike<'_> {
     fn range(&self) -> TextRange {
         match self {
@@ -561,6 +582,12 @@ impl<'a> From<&'a ast::FString> for StringLikePart<'a> {
 
 impl<'a> From<&StringLikePart<'a>> for AnyNodeRef<'a> {
     fn from(value: &StringLikePart<'a>) -> Self {
+        AnyNodeRef::from(*value)
+    }
+}
+
+impl<'a> From<StringLikePart<'a>> for AnyNodeRef<'a> {
+    fn from(value: StringLikePart<'a>) -> Self {
         match value {
             StringLikePart::String(part) => AnyNodeRef::StringLiteral(part),
             StringLikePart::Bytes(part) => AnyNodeRef::BytesLiteral(part),

crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/join_implicit_concatenated_string_assignment.py

@@ -291,3 +291,87 @@
     f"testeeeeeeeeeeeeeeeeeeeeeeeee{a
     =}" "moreeeeeeeeeeeeeeeeeetest"  # comment
 )
+
+
+# Trailing last-part comments
+
+a = (
+    "a"
+    "b"  # belongs to `b`
+)
+
+a: Literal[str] = (
+    "a"
+    "b"  # belongs to `b`
+)
+
+a += (
+    "a"
+    "b"  # belongs to `b`
+)
+
+a = (
+    r"a"
+    r"b"  # belongs to `b`
+)
+
+a = (
+    "a"
+    "b"
+)  # belongs to the assignment
+
+a = (((
+    "a"
+    "b"  # belongs to `b`
+)))
+
+a = (((
+    "a"
+    "b"
+) # belongs to the f-string expression
+))
+
+a = (
+    "a" "b"  # belongs to the f-string expression
+)
+
+a = (
+    "a" "b"
+    # belongs to the f-string expression
+)
+
+# There's no "right" answer if some parts are on the same line while others are on separate lines.
+# This is likely a comment for one of the last two parts but could also just be a comment for the entire f-string expression.
+# Because there's no right answer, follow what we do elsewhere and associate the comment with the outer-most node which
+# is the f-string expression.
+a = (
+    "a"
+    "b" "ccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc"  # belongs to the f-string expression
+)
+
+logger.error(
+    f"Failed to run task {task} for job"  
+    f"with id {str(job.id)}" # type: ignore[union-attr]
+)
+
+a = (10 +
+     "Exception in {call_back_name} "
+     f"'{msg}'"  # belongs to binary operation
+)
+
+a = 10 + (
+    "Exception in {call_back_name} "
+    f"'{msg}'"  # belongs to f-string
+)
+
+self._attr_unique_id = (
+    f"{self._device.temperature.group_address_state}_"
+    f"{self._device.target_temperature.group_address_state}_"
+    f"{self._device.target_temperature.group_address}_"
+    f"{self._device._setpoint_shift.group_address}"  # noqa: SLF001
+)
+
+return (
+    f"Exception in {call_back_name} when handling msg on "
+    f"'{msg.topic}': '{msg.payload}'"  # type: ignore[str-bytes-safe]
+)

crates/ruff_python_formatter/src/comments/placement.rs

@@ -1,19 +1,19 @@
-use std::cmp::Ordering;
-
 use ast::helpers::comment_indentation_after;
 use ruff_python_ast::whitespace::indentation;
 use ruff_python_ast::{
-    self as ast, AnyNodeRef, Comprehension, Expr, ModModule, Parameter, Parameters,
+    self as ast, AnyNodeRef, Comprehension, Expr, ModModule, Parameter, Parameters, StringLike,
 };
 use ruff_python_trivia::{
     find_only_token_in_range, first_non_trivia_token, indentation_at_offset, BackwardsTokenizer,
     CommentRanges, SimpleToken, SimpleTokenKind, SimpleTokenizer,
 };
 use ruff_source_file::LineRanges;
 use ruff_text_size::{Ranged, TextLen, TextRange};
+use std::cmp::Ordering;
 
 use crate::comments::visitor::{CommentPlacement, DecoratedComment};
 use crate::expression::expr_slice::{assign_comment_in_slice, ExprSliceCommentSection};
+use crate::expression::parentheses::is_expression_parenthesized;
 use crate::other::parameters::{
     assign_argument_separator_comment_placement, find_parameter_separators,
 };
@@ -355,6 +355,41 @@ fn handle_enclosed_comment<'a>(
         AnyNodeRef::ExprGenerator(generator) if generator.parenthesized => {
             handle_bracketed_end_of_line_comment(comment, source)
         }
+        AnyNodeRef::StmtReturn(_) => {
+            handle_trailing_implicit_concatenated_string_comment(comment, comment_ranges, source)
+        }
+        AnyNodeRef::StmtAssign(assignment)
+            if comment.preceding_node().is_some_and(|preceding| {
+                preceding.ptr_eq(AnyNodeRef::from(&*assignment.value))
+            }) =>
+        {
+            handle_trailing_implicit_concatenated_string_comment(comment, comment_ranges, source)
+        }
+        AnyNodeRef::StmtAnnAssign(assignment)
+            if comment.preceding_node().is_some_and(|preceding| {
+                assignment
+                    .value
+                    .as_deref()
+                    .is_some_and(|value| preceding.ptr_eq(value.into()))
+            }) =>
+        {
+            handle_trailing_implicit_concatenated_string_comment(comment, comment_ranges, source)
+        }
+        AnyNodeRef::StmtAugAssign(assignment)
+            if comment.preceding_node().is_some_and(|preceding| {
+                preceding.ptr_eq(AnyNodeRef::from(&*assignment.value))
+            }) =>
+        {
+            handle_trailing_implicit_concatenated_string_comment(comment, comment_ranges, source)
+        }
+        AnyNodeRef::StmtTypeAlias(assignment)
+            if comment.preceding_node().is_some_and(|preceding| {
+                preceding.ptr_eq(AnyNodeRef::from(&*assignment.value))
+            }) =>
+        {
+            handle_trailing_implicit_concatenated_string_comment(comment, comment_ranges, source)
+        }
+
         _ => CommentPlacement::Default(comment),
     }
 }
@@ -2086,6 +2121,75 @@ fn handle_comprehension_comment<'a>(
     CommentPlacement::Default(comment)
 }
 
+/// Handle end-of-line comments for parenthesized implicitly concatenated strings when used in
+/// a `FormatStatementLastExpression` context:
+///
+/// ```python
+/// a = (
+///     "a"
+///     "b"
+///     "c"  # comment
+/// )
+/// ```
+///
+/// `# comment` is a trailing comment of the last part and not a trailing comment of the entire f-string.
+/// Associating the comment with the last part is important or the assignment formatting might move
+/// the comment at the end of the assignment, making it impossible to suppress an error for the last part.
+///
+/// On the other hand, `# comment` is a trailing end-of-line f-string comment for:
+///
+/// ```python
+/// a = (
+///     "a" "b" "c"  # comment
+/// )
+///
+/// a = (
+///     "a"
+///     "b"
+///     "c"
+/// )  # comment
+/// ```
+///
+/// Associating the comment with the f-string is desired in those cases because it allows
+/// joining the string literals into a single string literal if it fits on the line.
+fn handle_trailing_implicit_concatenated_string_comment<'a>(
+    comment: DecoratedComment<'a>,
+    comment_ranges: &CommentRanges,
+    source: &str,
+) -> CommentPlacement<'a> {
+    if !comment.line_position().is_end_of_line() {
+        return CommentPlacement::Default(comment);
+    }
+
+    let Some(string_like) = comment
+        .preceding_node()
+        .and_then(|preceding| StringLike::try_from(preceding).ok())
+    else {
+        return CommentPlacement::Default(comment);
+    };
+
+    let mut parts = string_like.parts();
+
+    let (Some(last), Some(second_last)) = (parts.next_back(), parts.next_back()) else {
+        return CommentPlacement::Default(comment);
+    };
+
+    if source.contains_line_break(TextRange::new(second_last.end(), last.start()))
+        && is_expression_parenthesized(string_like.as_expression_ref(), comment_ranges, source)
+    {
+        let range = TextRange::new(last.end(), comment.start());
+
+        if !SimpleTokenizer::new(source, range)
+            .skip_trivia()
+            .any(|token| token.kind() == SimpleTokenKind::RParen)
+        {
+            return CommentPlacement::trailing(AnyNodeRef::from(last), comment);
+        }
+    }
+
+    CommentPlacement::Default(comment)
+}
+
 /// Returns `true` if the parameters are parenthesized (as in a function definition), or `false` if
 /// not (as in a lambda).
 fn are_parameters_parenthesized(parameters: &Parameters, contents: &str) -> bool {

crates/ruff_python_formatter/src/string/implicit.rs

@@ -99,7 +99,7 @@ impl Format<PyFormatContext<'_>> for FormatImplicitConcatenatedStringExpanded<'_
                 StringLikePart::FString(part) => part.format().fmt(f),
             });
 
-            let part_comments = comments.leading_dangling_trailing(&part);
+            let part_comments = comments.leading_dangling_trailing(part);
             joiner.entry(&format_args![
                 leading_comments(part_comments.leading),
                 format_part,

crates/ruff_python_formatter/tests/snapshots/format@expression__fstring.py.snap

@@ -1,7 +1,6 @@
 ---
 source: crates/ruff_python_formatter/tests/fixtures.rs
 input_file: crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/fstring.py
-snapshot_kind: text
 ---
 ## Input
 ```python