Generating return statements.

jedi/parser/pytree.py

@@ -76,9 +76,9 @@ def convert(grammar, raw_node):
         'import_from': pr.Import,
         'break_stmt': pr.KeywordStatement,
         'continue_stmt': pr.KeywordStatement,
-        'return_stmt': pr.KeywordStatement,
+        'return_stmt': pr.ReturnStmt,
         'raise_stmt': pr.KeywordStatement,
-        'yield_stmt': pr.KeywordStatement,
+        'yield_stmt': pr.ReturnStmt,
         'del_stmt': pr.KeywordStatement,
         'pass_stmt': pr.KeywordStatement,
         'global_stmt': pr.GlobalStmt,

jedi/parser/representation.py

@@ -397,7 +397,20 @@ class Scope(Simple, DocstringMixin):
     def returns(self):
         # Needed here for fast_parser, because the fast_parser splits and
         # returns will be in "normal" modules.
-        return [c for c in self.children if isinstance(c, ExprStmt)]
+        return self._search_in_scope(ReturnStmt)
+
+    def _search_in_scope(self, typ):
+        def scan(children):
+            elements = []
+            for element in children:
+                if isinstance(element, typ):
+                    elements.append(element)
+                elif is_node(element, 'suite') or is_node(element, 'simple_stmt'):
+                    elements += scan(element.children)
+            return elements
+
+        print('return', scan(self.children))
+        return scan(self.children)
 
     @property
     def statements(self):
@@ -972,6 +985,10 @@ class GlobalStmt(Simple):
         return self.children[1::2]
 
 
+class ReturnStmt(Simple):
+    pass
+
+
 class Statement(Simple, DocstringMixin):
     """
     This is the class for all the possible statements. Which means, this class