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@py_38__python38.py.snap

@@ -1,103 +0,0 @@
----
-source: crates/ruff_python_formatter/tests/fixtures.rs
-input_file: crates/ruff_python_formatter/resources/test/fixtures/black/py_38/python38.py
----
-## Input
-
-```py
-#!/usr/bin/env python3.8
-
-
-def starred_return():
-    my_list = ["value2", "value3"]
-    return "value1", *my_list
-
-
-def starred_yield():
-    my_list = ["value2", "value3"]
-    yield "value1", *my_list
-
-
-# all right hand side expressions allowed in regular assignments are now also allowed in
-# annotated assignments
-a : Tuple[ str, int] = "1", 2
-a: Tuple[int , ... ] = b,  *c, d
-def t():
-    a : str =  yield "a"
-```
-
-## Black Differences
-
-```diff
---- Black
-+++ Ruff
-@@ -8,7 +8,7 @@
- 
- def starred_yield():
-     my_list = ["value2", "value3"]
--    yield "value1", *my_list
-+    NOT_YET_IMPLEMENTED_ExprYield
- 
- 
- # all right hand side expressions allowed in regular assignments are now also allowed in
-@@ -18,4 +18,4 @@
- 
- 
- def t():
--    a: str = yield "a"
-+    a: str = NOT_YET_IMPLEMENTED_ExprYield
-```
-
-## Ruff Output
-
-```py
-#!/usr/bin/env python3.8
-
-
-def starred_return():
-    my_list = ["value2", "value3"]
-    return "value1", *my_list
-
-
-def starred_yield():
-    my_list = ["value2", "value3"]
-    NOT_YET_IMPLEMENTED_ExprYield
-
-
-# all right hand side expressions allowed in regular assignments are now also allowed in
-# annotated assignments
-a: Tuple[str, int] = "1", 2
-a: Tuple[int, ...] = b, *c, d
-
-
-def t():
-    a: str = NOT_YET_IMPLEMENTED_ExprYield
-```
-
-## Black Output
-
-```py
-#!/usr/bin/env python3.8
-
-
-def starred_return():
-    my_list = ["value2", "value3"]
-    return "value1", *my_list
-
-
-def starred_yield():
-    my_list = ["value2", "value3"]
-    yield "value1", *my_list
-
-
-# all right hand side expressions allowed in regular assignments are now also allowed in
-# annotated assignments
-a: Tuple[str, int] = "1", 2
-a: Tuple[int, ...] = b, *c, d
-
-
-def t():
-    a: str = yield "a"
-```
-
-

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():

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

@@ -311,16 +311,15 @@ d={'a':1,
  
  def yield_expr():
      # fmt: off
--    yield hello
+     yield hello
 -    'unformatted'
-+    NOT_YET_IMPLEMENTED_ExprYield
 +    "unformatted"
      # fmt: on
      "formatted"
      # fmt: off
 -    ( yield hello )
 -    'unformatted'
-+    (NOT_YET_IMPLEMENTED_ExprYield)
++    (yield hello)
 +    "unformatted"
      # fmt: on
  
@@ -379,7 +378,7 @@ d={'a':1,
  )
  # fmt: off
 -yield  'hello'
-+NOT_YET_IMPLEMENTED_ExprYield
++yield "hello"
  # No formatting to the end of the file
 -l=[1,2,3]
 -d={'a':1,
@@ -500,12 +499,12 @@ def testlist_star_expr():
 
 def yield_expr():
     # fmt: off
-    NOT_YET_IMPLEMENTED_ExprYield
+    yield hello
     "unformatted"
     # fmt: on
     "formatted"
     # fmt: off
-    (NOT_YET_IMPLEMENTED_ExprYield)
+    (yield hello)
     "unformatted"
     # fmt: on
 
@@ -617,7 +616,7 @@ cfg.rule(
     xxxxxxxxxx_xxxxxxxxxxx_xxxxxxx_xxxxxxxxx=5,
 )
 # fmt: off
-NOT_YET_IMPLEMENTED_ExprYield
+yield "hello"
 # No formatting to the end of the file
 l = [1, 2, 3]
 d = {"a": 1, "b": 2}

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

@@ -141,23 +141,6 @@ def __await__(): return (yield)
  
  
  def long_lines():
-@@ -135,14 +130,8 @@
-     a,
-     **kwargs,
- ) -> A:
--    return (
--        yield from A(
--            very_long_argument_name1=very_long_value_for_the_argument,
--            very_long_argument_name2=very_long_value_for_the_argument,
--            **kwargs,
--        )
--    )
-+    return NOT_YET_IMPLEMENTED_ExprYieldFrom
- 
- 
- def __await__():
--    return (yield)
-+    return NOT_YET_IMPLEMENTED_ExprYield
 ```
 
 ## Ruff Output
@@ -295,11 +278,17 @@ def f(
     a,
     **kwargs,
 ) -> A:
-    return NOT_YET_IMPLEMENTED_ExprYieldFrom
+    return (
+        yield from A(
+            very_long_argument_name1=very_long_value_for_the_argument,
+            very_long_argument_name2=very_long_value_for_the_argument,
+            **kwargs,
+        )
+    )
 
 
 def __await__():
-    return NOT_YET_IMPLEMENTED_ExprYield
+    return (yield)
 ```
 
 ## Black Output

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

@@ -105,21 +105,6 @@ async def main():
  
  
  async def main():
-@@ -78,7 +81,7 @@
- 
- 
- async def main():
--    await (yield x)
-+    await NOT_YET_IMPLEMENTED_ExprYield
- 
- 
- async def main():
-@@ -90,4 +93,4 @@
- 
- 
- async def main():
--    await (yield)
-+    await NOT_YET_IMPLEMENTED_ExprYield
 ```
 
 ## Ruff Output
@@ -208,7 +193,7 @@ async def main():
 
 
 async def main():
-    await NOT_YET_IMPLEMENTED_ExprYield
+    await (yield x)
 
 
 async def main():
@@ -220,7 +205,7 @@ async def main():
 
 
 async def main():
-    await NOT_YET_IMPLEMENTED_ExprYield
+    await (yield)
 ```
 
 ## Black Output

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

@@ -0,0 +1,121 @@
+---
+source: crates/ruff_python_formatter/tests/fixtures.rs
+input_file: crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/yield.py
+---
+## Input
+```py
+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)
+
+```
+
+## Output
+```py
+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/tests/snapshots/format@expression__yield_from.py.snap

@@ -0,0 +1,74 @@
+---
+source: crates/ruff_python_formatter/tests/fixtures.rs
+input_file: crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/yield_from.py
+---
+## Input
+```py
+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
+    )
+
+```
+
+## Output
+```py
+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)
+```
+
+
+