Format ExprYield/ExprYieldFrom (#5921)

Co-authored-by: Micha Reiser <micha@reiser.io>

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

@@ -0,0 +1,59 @@
+l = [1,2,3,4]
+
+def foo():
+    yield l
+
+    b = yield l
+
+    c = [
+        (yield l) , (
+        yield l
+
+    )]
+
+    with (
+        # Some comment
+        yield
+    ):
+        pass
+
+    if (yield):
+        # comment
+        pass
+
+
+    (yield a, b) = (1, 2)
+
+    # some comment
+    for e in l : yield e # some comment
+
+    for e in l:
+
+
+        # some comment
+
+        yield e
+
+        # trail comment
+
+    for e in l:
+        # comment
+        yield (((((((e)))))))  # Too many parentheses
+        # comment
+
+
+    for ridiculouslylongelementnameeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee in l:
+        yield ridiculouslylongelementnameeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee
+
+
+    for x in l: #comment
+        yield x + (2 * 4) # trailing comment
+
+    while (
+
+        yield l
+    ):
+        pass
+
+    yield from (yield l)
+

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

@@ -0,0 +1,35 @@
+l = [1,2,3,4]
+
+
+def foo():
+    yield from l # some comment
+
+    # weird indents
+    yield\
+                    from\
+                        l
+                    # indented trailing comment
+
+    a = yield from l
+
+    with (
+        # Comment
+        yield from l
+        # Comment
+    ):
+        pass
+
+    c = [(yield from l) , (
+        yield from l
+
+    )]
+
+    while (
+        yield from l
+    ):
+        pass
+
+    yield (
+        yield from l
+    )
+

crates/ruff_python_formatter/src/expression/expr_yield.rs

@@ -1,6 +1,8 @@
 use crate::context::PyFormatContext;
-use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses};
-use crate::{not_yet_implemented, FormatNodeRule, PyFormatter};
+use crate::expression::maybe_parenthesize_expression;
+use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses, Parenthesize};
+use crate::{FormatNodeRule, PyFormatter};
+use ruff_formatter::prelude::{space, text};
 use ruff_formatter::{write, Buffer, FormatResult};
 use ruff_python_ast::node::AnyNodeRef;
 use rustpython_parser::ast::ExprYield;
@@ -10,16 +12,40 @@ pub struct FormatExprYield;
 
 impl FormatNodeRule<ExprYield> for FormatExprYield {
     fn fmt_fields(&self, item: &ExprYield, f: &mut PyFormatter) -> FormatResult<()> {
-        write!(f, [not_yet_implemented(item)])
+        let ExprYield { range: _, value } = item;
+
+        if let Some(val) = value {
+            write!(
+                f,
+                [
+                    text("yield"),
+                    space(),
+                    maybe_parenthesize_expression(val, item, Parenthesize::IfRequired)
+                ]
+            )?;
+        } else {
+            write!(f, [&text("yield")])?;
+        }
+        Ok(())
     }
 }
 
 impl NeedsParentheses for ExprYield {
     fn needs_parentheses(
         &self,
-        _parent: AnyNodeRef,
+        parent: AnyNodeRef,
         _context: &PyFormatContext,
     ) -> OptionalParentheses {
-        OptionalParentheses::Multiline
+        // 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:
+        //  1. Right hand side of an assignment, e.g. `x = yield y`
+        //  2. Yield statement, e.g. `def foo(): yield y`
+        // 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
+        } else {
+            OptionalParentheses::Always
+        }
     }
 }

crates/ruff_python_formatter/src/expression/expr_yield_from.rs

@@ -1,6 +1,8 @@
 use crate::context::PyFormatContext;
-use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses};
-use crate::{not_yet_implemented, FormatNodeRule, PyFormatter};
+use crate::expression::maybe_parenthesize_expression;
+use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses, Parenthesize};
+use crate::{FormatNodeRule, PyFormatter};
+use ruff_formatter::prelude::{space, text};
 use ruff_formatter::{write, Buffer, FormatResult};
 use ruff_python_ast::node::AnyNodeRef;
 use rustpython_parser::ast::ExprYieldFrom;
@@ -10,16 +12,37 @@ pub struct FormatExprYieldFrom;
 
 impl FormatNodeRule<ExprYieldFrom> for FormatExprYieldFrom {
     fn fmt_fields(&self, item: &ExprYieldFrom, f: &mut PyFormatter) -> FormatResult<()> {
-        write!(f, [not_yet_implemented(item)])
+        let ExprYieldFrom { range: _, value } = item;
+
+        write!(
+            f,
+            [
+                text("yield from"),
+                space(),
+                maybe_parenthesize_expression(value, item, Parenthesize::IfRequired)
+            ]
+        )?;
+
+        Ok(())
     }
 }
 
 impl NeedsParentheses for ExprYieldFrom {
     fn needs_parentheses(
         &self,
-        _parent: AnyNodeRef,
+        parent: AnyNodeRef,
         _context: &PyFormatContext,
     ) -> OptionalParentheses {
-        OptionalParentheses::Multiline
+        // 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:
+        //  1. Right hand side of an assignment, e.g. `x = yield y`
+        //  2. Yield statement, e.g. `def foo(): yield y`
+        // 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
+        } else {
+            OptionalParentheses::Always
+        }
     }
 }

crates/ruff_python_formatter/tests/snapshots/black_compatibility@miscellaneous__debug_visitor.py.snap

@@ -56,8 +56,7 @@ class DebugVisitor(Visitor[T]):
 +            out(f"NOT_YET_IMPLEMENTED_ExprJoinedStr", fg="yellow")
              self.tree_depth += 1
              for child in node.children:
--                yield from self.visit(child)
-+                NOT_YET_IMPLEMENTED_ExprYieldFrom
+                 yield from self.visit(child)
  
              self.tree_depth -= 1
 -            out(f'{indent}/{_type}', fg='yellow', bold=False)
@@ -97,7 +96,7 @@ class DebugVisitor(Visitor[T]):
             out(f"NOT_YET_IMPLEMENTED_ExprJoinedStr", fg="yellow")
             self.tree_depth += 1
             for child in node.children:
-                NOT_YET_IMPLEMENTED_ExprYieldFrom
+                yield from self.visit(child)
 
             self.tree_depth -= 1
             out(f"NOT_YET_IMPLEMENTED_ExprJoinedStr", fg="yellow", bold=False)

crates/ruff_python_formatter/tests/snapshots/black_compatibility@miscellaneous__long_strings_flag_disabled.py.snap

@@ -343,12 +343,8 @@ long_unmergable_string_with_pragma = (
  )
 
  func_with_bad_parens_that_wont_fit_in_one_line(
-@@ -271,10 +267,10 @@
-
-
- def foo():
--    yield "This is a really long string that can't possibly be expected to fit all together on one line. In fact it may even take up three or more lines... like four or five... but probably just three."
-+    NOT_YET_IMPLEMENTED_ExprYield
+@@ -274,7 +270,7 @@
+     yield "This is a really long string that can't possibly be expected to fit all together on one line. In fact it may even take up three or more lines... like four or five... but probably just three."
 
 
 -x = f"This is a {{really}} long string that needs to be split without a doubt (i.e. most definitely). In short, this {string} that can't possibly be {{expected}} to fit all together on one line. In {fact} it may even take up three or more lines... like four or five... but probably just four."
@@ -630,7 +626,7 @@ raw_strings = r"Don't" " get" r" merged" " unless they are all raw."
 
 
 def foo():
-    NOT_YET_IMPLEMENTED_ExprYield
+    yield "This is a really long string that can't possibly be expected to fit all together on one line. In fact it may even take up three or more lines... like four or five... but probably just three."
 
 
 x = f"NOT_YET_IMPLEMENTED_ExprJoinedStr"

crates/ruff_python_formatter/tests/snapshots/black_compatibility@simple_cases__expression.py.snap

@@ -357,21 +357,6 @@ last_call()
  Ø = set()
  authors.łukasz.say_thanks()
  mapping = {
-@@ -237,10 +234,10 @@
- 
- 
- def gen():
--    yield from outside_of_generator
--    a = yield
--    b = yield
--    c = yield
-+    NOT_YET_IMPLEMENTED_ExprYieldFrom
-+    a = NOT_YET_IMPLEMENTED_ExprYield
-+    b = NOT_YET_IMPLEMENTED_ExprYield
-+    c = NOT_YET_IMPLEMENTED_ExprYield
- 
- 
- async def f():
 @@ -328,13 +325,18 @@
  ):
      return True
@@ -645,10 +630,10 @@ mapping = {
 
 
 def gen():
-    NOT_YET_IMPLEMENTED_ExprYieldFrom
-    a = NOT_YET_IMPLEMENTED_ExprYield
-    b = NOT_YET_IMPLEMENTED_ExprYield
-    c = NOT_YET_IMPLEMENTED_ExprYield
+    yield from outside_of_generator
+    a = yield
+    b = yield
+    c = yield
 
 
 async def f():