From a21ec2c0ad80056e44102bc8c988107ac67c9902 Mon Sep 17 00:00:00 2001
From: Dave Halter <davidhalter88@gmail.com>
Date: Fri, 1 Sep 2017 01:06:07 +0200
Subject: [PATCH] Better yield/yield from support.

---
 parso/python/tree.py     | 17 +++++++++++++++--
 test/test_parser_tree.py | 17 +++++++++++++++++
 2 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/parso/python/tree.py b/parso/python/tree.py
index a9121b7..11ebb2e 100644
--- a/parso/python/tree.py
+++ b/parso/python/tree.py
@@ -519,8 +519,21 @@ class Function(ClassOrFunc):
         """
         Returns a generator of `yield_expr`.
         """
-        # TODO This is incorrect, yields are also possible in a statement.
-        return self._search_in_scope('yield_expr')
+        def scan(children):
+            for element in children:
+                try:
+                    nested_children = element.children
+                except AttributeError:
+                    if element.value == 'yield':
+                        if element.parent.type == 'yield_expr':
+                            yield element.parent
+                        else:
+                            yield element
+                else:
+                    for result in scan(nested_children):
+                        yield result
+
+        return scan(self.children)
 
     def iter_return_stmts(self):
         """
diff --git a/test/test_parser_tree.py b/test/test_parser_tree.py
index ca7abb5..c0f3d96 100644
--- a/test/test_parser_tree.py
+++ b/test/test_parser_tree.py
@@ -115,3 +115,20 @@ def test_ellipsis_py2(each_py2_version):
     subscript = trailer.children[1]
     assert subscript.type == 'subscript'
     assert [leaf.value for leaf in subscript.children] == ['.', '.', '.']
+
+
+def get_yield_exprs(code, version):
+    return list(parse(code, version=version).children[0].iter_yield_exprs())
+
+
+def test_yields(each_version):
+    y, = get_yield_exprs('def x(): yield', each_version)
+    assert y.value == 'yield'
+
+    y, = get_yield_exprs('def x(): (yield 1)', each_version)
+    assert y.type == 'yield_expr'
+
+
+def test_yield_from():
+    y, = get_yield_exprs('def x(): (yield from 1)', '3.3')
+    assert y.type == 'yield_expr'