Use optional parentheses for tuples in return statement

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

@@ -73,6 +73,18 @@ def foo():
 
     yield aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb + ccccccccccccccccccccccccccccccccccccccccccccccccccccccc
 
+    yield len(self.nodeseeeeeeeee), sum(
+        len(node.parents) for node in self.node_map.values()
+    )
+
+    a = (yield len(self.nodeseeeeeeeee), sum(
+        len(node.parents) for node in self.node_map.values()
+    ))
+
+    a = yield (yield len(self.nodeseeeeeeeee), sum(
+        len(node.parents) for node in self.node_map.values()
+    ))
+
 yield ("Cache key will cause errors if used with memcached: %r " "(longer than %s)" % (
     key,
     MEMCACHE_MAX_KEY_LENGTH,

crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/return.py

@@ -0,0 +1,16 @@
+
+return len(self.nodeseeeeeeeee), sum(
+    len(node.parents) for node in self.node_map.values()
+)
+
+
+return len(self.nodeseeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee), sum(
+    len(node.parents) for node in self.node_map.values()
+)
+
+
+return (
+    len(self.nodeseeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee), sum(
+    len(node.parents) for node in self.node_map.values()
+    )
+)

crates/ruff_python_formatter/src/expression/expr_tuple.rs

@@ -32,6 +32,20 @@ pub enum TupleParentheses {
     /// ```
     Preserve,
 
+    /// The same as [`Self::Default`] except that it uses [`optional_parentheses`] rather than
+    /// [`parenthesize_if_expands`]. This avoids adding parentheses if breaking any containing parenthesized
+    /// expression makes the tuple fit.
+    ///
+    /// Avoids adding parentheses around the tuple because breaking the `sum` call expression is sufficient
+    /// to make it fit.
+    ///
+    /// ```python
+    /// return len(self.nodeseeeeeeeee), sum(
+    //     len(node.parents) for node in self.node_map.values()
+    // )
+    /// ```
+    OptionalParentheses,
+
     /// Handle the special cases where we don't include parentheses at all.
     ///
     /// Black never formats tuple targets of for loops with parentheses if inside a comprehension.
@@ -159,7 +173,7 @@ impl FormatNodeRule<ExprTuple> for FormatExprTuple {
                         .finish()
                 }
                 TupleParentheses::Preserve => group(&ExprSequence::new(item)).fmt(f),
-                TupleParentheses::NeverPreserve => {
+                TupleParentheses::NeverPreserve | TupleParentheses::OptionalParentheses => {
                     optional_parentheses(&ExprSequence::new(item)).fmt(f)
                 }
                 TupleParentheses::Default => {

crates/ruff_python_formatter/src/expression/expr_yield.rs

@@ -1,11 +1,15 @@
-use crate::context::PyFormatContext;
+use crate::context::{NodeLevel, PyFormatContext};
+use crate::expression::expr_tuple::TupleParentheses;
 use crate::expression::maybe_parenthesize_expression;
-use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses, Parenthesize};
+use crate::expression::parentheses::{
+    is_expression_parenthesized, NeedsParentheses, OptionalParentheses, Parenthesize,
+};
+use crate::pattern::pattern_match_sequence::SequenceType::Tuple;
 use crate::prelude::*;
 use crate::{FormatNodeRule, PyFormatter};
 use ruff_formatter::write;
 use ruff_python_ast::node::AnyNodeRef;
-use ruff_python_ast::{Expr, ExprYield, ExprYieldFrom, Ranged};
+use ruff_python_ast::{Expr, ExprYield, ExprYieldFrom, ExpressionRef, Ranged};
 use ruff_text_size::TextRange;
 
 pub(super) enum AnyExpressionYield<'a> {
@@ -39,7 +43,7 @@ impl NeedsParentheses for AnyExpressionYield<'_> {
     fn needs_parentheses(
         &self,
         parent: AnyNodeRef,
-        _context: &PyFormatContext,
+        context: &PyFormatContext,
     ) -> OptionalParentheses {
         // According to https://docs.python.org/3/reference/grammar.html There are two situations
         // where we do not want to always parenthesize a yield expression:
@@ -48,7 +52,15 @@ impl NeedsParentheses for AnyExpressionYield<'_> {
         // We catch situation 1 below. Situation 2 does not need to be handled here as
         // FormatStmtExpr, does not add parenthesis
         if parent.is_stmt_assign() || parent.is_stmt_ann_assign() || parent.is_stmt_aug_assign() {
-            OptionalParentheses::Multiline
+            if self
+                .value()
+                .and_then(|value| value.as_tuple_expr())
+                .is_some_and(|tuple| !context.comments().has_leading(tuple))
+            {
+                OptionalParentheses::Never
+            } else {
+                OptionalParentheses::Multiline
+            }
         } else {
             OptionalParentheses::Always
         }
@@ -85,19 +97,24 @@ impl Format<PyFormatContext<'_>> for AnyExpressionYield<'_> {
         };
 
         if let Some(val) = self.value() {
-            write!(
-                f,
-                [
-                    text(keyword),
-                    space(),
-                    maybe_parenthesize_expression(val, self, Parenthesize::Optional)
-                ]
-            )?;
+            write!(f, [text(keyword), space()])?;
+
+            match val {
+                Expr::Tuple(tuple)
+                    if !f.context().comments().has_leading(tuple)
+                        && !f.context().node_level().is_parenthesized() =>
+                {
+                    tuple
+                        .format()
+                        .with_options(TupleParentheses::OptionalParentheses)
+                        .fmt(f)
+                }
+                _ => maybe_parenthesize_expression(val, self, Parenthesize::Optional).fmt(f),
+            }
         } else {
             // ExprYieldFrom always has Some(value) so we should never get a bare `yield from`
-            write!(f, [&text(keyword)])?;
+            text(keyword).fmt(f)
         }
-        Ok(())
     }
 }
 

crates/ruff_python_formatter/src/statement/stmt_return.rs

@@ -1,7 +1,8 @@
 use ruff_formatter::write;
-use ruff_python_ast::StmtReturn;
+use ruff_python_ast::{Expr, StmtReturn};
 
 use crate::comments::{SourceComment, SuppressionKind};
+use crate::expression::expr_tuple::TupleParentheses;
 use crate::expression::maybe_parenthesize_expression;
 use crate::expression::parentheses::Parenthesize;
 use crate::prelude::*;
@@ -13,17 +14,31 @@ pub struct FormatStmtReturn;
 impl FormatNodeRule<StmtReturn> for FormatStmtReturn {
     fn fmt_fields(&self, item: &StmtReturn, f: &mut PyFormatter) -> FormatResult<()> {
         let StmtReturn { range: _, value } = item;
-        if let Some(value) = value {
-            write!(
-                f,
-                [
-                    text("return"),
-                    space(),
-                    maybe_parenthesize_expression(value, item, Parenthesize::IfBreaks)
-                ]
-            )
-        } else {
-            text("return").fmt(f)
+
+        text("return").fmt(f)?;
+
+        match value.as_deref() {
+            Some(Expr::Tuple(tuple)) if !f.context().comments().has_leading(tuple) => {
+                write!(
+                    f,
+                    [
+                        space(),
+                        tuple
+                            .format()
+                            .with_options(TupleParentheses::OptionalParentheses)
+                    ]
+                )
+            }
+            Some(value) => {
+                write!(
+                    f,
+                    [
+                        space(),
+                        maybe_parenthesize_expression(value, item, Parenthesize::IfBreaks)
+                    ]
+                )
+            }
+            None => Ok(()),
         }
     }
 

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

@@ -252,33 +252,23 @@ self.assertEqual(
 
 def test():
     return (
-        (
-            "((TIME_TO_SEC(%(lhs)s) * 1000000 + MICROSECOND(%(lhs)s)) -"
-            " (TIME_TO_SEC(%(rhs)s) * 1000000 + MICROSECOND(%(rhs)s)))"
-        ) % {"lhs": lhs_sql, "rhs": rhs_sql},
-        tuple(lhs_params) * 2 + tuple(rhs_params) * 2,
-    )
+        "((TIME_TO_SEC(%(lhs)s) * 1000000 + MICROSECOND(%(lhs)s)) -"
+        " (TIME_TO_SEC(%(rhs)s) * 1000000 + MICROSECOND(%(rhs)s)))"
+    ) % {"lhs": lhs_sql, "rhs": rhs_sql}, tuple(lhs_params) * 2 + tuple(rhs_params) * 2
 
 
 def test2():
-    return (
-        "RETURNING %s INTO %s"
-        % (
-            ", ".join(field_names),
-            ", ".join(["%s"] * len(params)),
-        ),
-        tuple(params),
-    )
+    return "RETURNING %s INTO %s" % (
+        ", ".join(field_names),
+        ", ".join(["%s"] * len(params)),
+    ), tuple(params)
 
 
 def test3():
     return (
-        (
-            "(CASE WHEN JSON_TYPE(%s, %%s) IN (%s) "
-            "THEN JSON_TYPE(%s, %%s) ELSE JSON_EXTRACT(%s, %%s) END)"
-        ) % (lhs, datatype_values, lhs, lhs),
-        (tuple(params) + (json_path,)) * 3,
-    )
+        "(CASE WHEN JSON_TYPE(%s, %%s) IN (%s) "
+        "THEN JSON_TYPE(%s, %%s) ELSE JSON_EXTRACT(%s, %%s) END)"
+    ) % (lhs, datatype_values, lhs, lhs), (tuple(params) + (json_path,)) * 3
 ```
 
 

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

@@ -79,6 +79,18 @@ def foo():
 
     yield aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa + bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb + ccccccccccccccccccccccccccccccccccccccccccccccccccccccc
 
+    yield len(self.nodeseeeeeeeee), sum(
+        len(node.parents) for node in self.node_map.values()
+    )
+
+    a = (yield len(self.nodeseeeeeeeee), sum(
+        len(node.parents) for node in self.node_map.values()
+    ))
+
+    a = yield (yield len(self.nodeseeeeeeeee), sum(
+        len(node.parents) for node in self.node_map.values()
+    ))
+
 yield ("Cache key will cause errors if used with memcached: %r " "(longer than %s)" % (
     key,
     MEMCACHE_MAX_KEY_LENGTH,
@@ -186,6 +198,21 @@ def foo():
         + ccccccccccccccccccccccccccccccccccccccccccccccccccccccc
     )
 
+    yield len(self.nodeseeeeeeeee), sum(
+        len(node.parents) for node in self.node_map.values()
+    )
+
+    a = yield len(self.nodeseeeeeeeee), sum(
+        len(node.parents) for node in self.node_map.values()
+    )
+
+    a = yield (
+        yield (
+            len(self.nodeseeeeeeeee),
+            sum(len(node.parents) for node in self.node_map.values()),
+        )
+    )
+
 
 yield (
     "Cache key will cause errors if used with memcached: %r "

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

@@ -0,0 +1,44 @@
+---
+source: crates/ruff_python_formatter/tests/fixtures.rs
+input_file: crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/return.py
+---
+## Input
+```py
+
+return len(self.nodeseeeeeeeee), sum(
+    len(node.parents) for node in self.node_map.values()
+)
+
+
+return len(self.nodeseeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee), sum(
+    len(node.parents) for node in self.node_map.values()
+)
+
+
+return (
+    len(self.nodeseeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee), sum(
+    len(node.parents) for node in self.node_map.values()
+    )
+)
+```
+
+## Output
+```py
+return len(self.nodeseeeeeeeee), sum(
+    len(node.parents) for node in self.node_map.values()
+)
+
+
+return len(
+    self.nodeseeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee
+), sum(len(node.parents) for node in self.node_map.values())
+
+
+return (
+    len(self.nodeseeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee),
+    sum(len(node.parents) for node in self.node_map.values()),
+)
+```
+
+
+