Implement yield-to-yield from conversion (#1544)

README.md

@@ -693,6 +693,7 @@ For more, see [pyupgrade](https://pypi.org/project/pyupgrade/3.2.0/) on PyPI.
 | UP025 | RewriteUnicodeLiteral | Remove unicode literals from strings | 🛠 |
 | UP026 | RewriteMockImport | `mock` is deprecated, use `unittest.mock` | 🛠 |
 | UP027 | RewriteListComprehension | Replace unpacked list comprehension with a generator expression | 🛠 |
+| UP028 | RewriteYieldFrom | Replace `yield` over `for` loop with `yield from` | 🛠 |
 
 ### pep8-naming (N)
 

resources/test/fixtures/pyupgrade/UP028_0.py

@@ -0,0 +1,74 @@
+def f():
+    for x in y:
+        yield x
+
+
+def g():
+    for x, y in z:
+        yield (x, y)
+
+
+def h():
+    for x in [1, 2, 3]:
+        yield x
+
+
+def i():
+    for x in {x for x in y}:
+        yield x
+
+
+def j():
+    for x in (1, 2, 3):
+        yield x
+
+
+def k():
+    for x, y in {3: "x", 6: "y"}:
+        yield x, y
+
+
+def f():  # Comment one\n'
+    # Comment two\n'
+    for x, y in {  # Comment three\n'
+        3: "x",  # Comment four\n'
+        # Comment five\n'
+        6: "y",  # Comment six\n'
+    }:  # Comment seven\n'
+        # Comment eight\n'
+        yield x, y  # Comment nine\n'
+        # Comment ten',
+
+
+def f():
+    for x, y in [{3: (3, [44, "long ss"]), 6: "y"}]:
+        yield x, y
+
+
+def f():
+    for x, y in z():
+        yield x, y
+
+def f():
+    def func():
+        # This comment is preserved\n'
+        for x, y in z():  # Comment one\n'
+            # Comment two\n'
+            yield x, y  # Comment three\n'
+            # Comment four\n'
+# Comment\n'
+def g():
+    print(3)
+
+
+def f():
+    for x in y:
+        yield x
+    for z in x:
+        yield z
+
+
+def f():
+    for x, y in z():
+        yield x, y
+    x = 1

resources/test/fixtures/pyupgrade/UP028_1.py

@@ -0,0 +1,113 @@
+# These should NOT change
+def f():
+    for x in z:
+        yield
+
+
+def f():
+    for x in z:
+        yield y
+
+
+def f():
+    for x, y in z:
+        yield x
+
+
+def f():
+    for x, y in z:
+        yield y
+
+
+def f():
+    for a, b in z:
+        yield x, y
+
+
+def f():
+    for x, y in z:
+        yield y, x
+
+
+def f():
+    for x, y, c in z:
+        yield x, y
+
+
+def f():
+    for x in z:
+        x = 22
+        yield x
+
+
+def f():
+    for x in z:
+        yield x
+    else:
+        print("boom!")
+
+
+def f():
+    for x in range(5):
+        yield x
+    print(x)
+
+
+def f():
+    def g():
+        print(x)
+
+    for x in range(5):
+        yield x
+    g()
+
+
+def f():
+    def g():
+        def h():
+            print(x)
+
+        return h
+
+    for x in range(5):
+        yield x
+    g()()
+
+
+def f(x):
+    for x in y:
+        yield x
+    del x
+
+
+async def f():
+    for x in y:
+        yield x
+
+
+def f():
+    x = 1
+    print(x)
+    for x in y:
+        yield x
+
+
+def f():
+    for x in y:
+        yield x
+    print(x)
+
+
+def f():
+    for x in y:
+        yield x
+    z = lambda: x
+
+
+def f():
+    for x in y:
+        yield x
+
+    class C:
+        def __init__(self):
+            print(x)

ruff.schema.json

@@ -969,6 +969,7 @@
         "UP025",
         "UP026",
         "UP027",
+        "UP028",
         "W",
         "W2",
         "W29",

src/checkers/ast.rs

@@ -1357,6 +1357,9 @@ where
                     body,
                     &Documentable::Function,
                 );
+                if self.settings.enabled.contains(&CheckCode::UP028) {
+                    pyupgrade::plugins::rewrite_yield_from(self, stmt);
+                }
                 let scope = transition_scope(&self.visible_scope, stmt, &Documentable::Function);
                 self.definitions
                     .push((definition, scope.visibility.clone()));

src/pyupgrade/mod.rs

@@ -48,6 +48,8 @@ mod tests {
     #[test_case(CheckCode::UP025, Path::new("UP025.py"); "UP025")]
     #[test_case(CheckCode::UP026, Path::new("UP026.py"); "UP026")]
     #[test_case(CheckCode::UP027, Path::new("UP027.py"); "UP027")]
+    #[test_case(CheckCode::UP028, Path::new("UP028_0.py"); "UP028_0")]
+    #[test_case(CheckCode::UP028, Path::new("UP028_1.py"); "UP028_1")]
     fn checks(check_code: CheckCode, path: &Path) -> Result<()> {
         let snapshot = format!("{}_{}", check_code.as_ref(), path.to_string_lossy());
         let checks = test_path(

src/pyupgrade/plugins/mod.rs

@@ -12,6 +12,7 @@ pub use replace_universal_newlines::replace_universal_newlines;
 pub use rewrite_c_element_tree::replace_c_element_tree;
 pub use rewrite_mock_import::{rewrite_mock_attribute, rewrite_mock_import};
 pub use rewrite_unicode_literal::rewrite_unicode_literal;
+pub use rewrite_yield_from::rewrite_yield_from;
 pub use super_call_with_parameters::super_call_with_parameters;
 pub use type_of_primitive::type_of_primitive;
 pub use typing_text_str_alias::typing_text_str_alias;
@@ -38,6 +39,7 @@ mod replace_universal_newlines;
 mod rewrite_c_element_tree;
 mod rewrite_mock_import;
 mod rewrite_unicode_literal;
+mod rewrite_yield_from;
 mod super_call_with_parameters;
 mod type_of_primitive;
 mod typing_text_str_alias;