diff --git a/jedi/parser/pytree.py b/jedi/parser/pytree.py index 1f31ed37..1cc0f0de 100644 --- a/jedi/parser/pytree.py +++ b/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, diff --git a/jedi/parser/representation.py b/jedi/parser/representation.py index 51a62dab..4f6abf6b 100644 --- a/jedi/parser/representation.py +++ b/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