🐛 Fix yield_from_generator being too broad

octoprint_codemods/yield_from_generator.py

@@ -1,4 +1,4 @@
-from typing import Union
+from typing import Union, cast
 
 import libcst as cst
 import libcst.matchers as m
@@ -20,14 +20,29 @@ def leave_For(
         self, original_node: cst.For, updated_node: cst.For
     ) -> Union[cst.For, cst.SimpleStatementLine]:
         if m.matches(
-            updated_node.body,
-            m.IndentedBlock(body=[m.SimpleStatementLine(body=[m.Expr(value=m.Yield())])]),
+            updated_node,
+            m.For(
+                target=m.Name(),
+                body=m.IndentedBlock(
+                    body=[m.SimpleStatementLine(body=[m.Expr(value=m.Yield(m.Name()))])]
+                ),
+            ),
         ):
-            self._report_node(original_node)
-            self.count += 1
-            updated_node = cst.SimpleStatementLine(
-                body=[cst.Expr(value=cst.Yield(value=cst.From(item=updated_node.iter)))]
-            )
+            target = updated_node.target.value
+            block = cast(cst.IndentedBlock, updated_node.body)
+            simple_stmt = cast(cst.SimpleStatementLine, block.body[0])
+            expr_stmt = cast(cst.Expr, simple_stmt.body[0])
+            yield_stmt = cast(cst.Yield, expr_stmt.value)
+            yielded = cast(cst.Name, yield_stmt.value).value
+
+            if target == yielded:
+                self._report_node(original_node)
+                self.count += 1
+                updated_node = cst.SimpleStatementLine(
+                    body=[
+                        cst.Expr(value=cst.Yield(value=cst.From(item=updated_node.iter)))
+                    ]
+                )
         return updated_node
 
 

tests/expected/yield_from_generator.py

@@ -1,6 +1,10 @@
 data = (x for x in range(10))
 
 
+def fnord(x):
+    return x
+
+
 def a():
     yield from data
 
@@ -11,3 +15,19 @@ def b():
 
 def c():
     yield from data
+
+
+def d():
+    for x in data:
+        yield fnord(x)
+
+
+def e():
+    for x in data:
+        yield x, True
+
+
+def f():
+    l = "aaaaaaaaaaaaaaaaaaa"
+    for x in range(3):
+        yield l[x:]

tests/fixtures/yield_from_generator.py

@@ -1,6 +1,10 @@
 data = (x for x in range(10))
 
 
+def fnord(x):
+    return x
+
+
 def a():
     for entry in data:
         yield entry
@@ -13,3 +17,19 @@ def b():
 
 def c():
     yield from data
+
+
+def d():
+    for x in data:
+        yield fnord(x)
+
+
+def e():
+    for x in data:
+        yield x, True
+
+
+def f():
+    l = "aaaaaaaaaaaaaaaaaaa"
+    for x in range(3):
+        yield l[x:]