diff --git a/jedi/evaluate/representation.py b/jedi/evaluate/representation.py index d0a19fd6..c0a9b8fa 100644 --- a/jedi/evaluate/representation.py +++ b/jedi/evaluate/representation.py @@ -319,7 +319,7 @@ class FunctionExecutionContext(context.TreeContext): if check_yields: types = set() - returns = funcdef.yields + returns = funcdef.iter_yield_exprs() else: returns = funcdef.iter_return_stmts() types = set(docstrings.infer_return_types(self.function_context)) @@ -352,7 +352,7 @@ class FunctionExecutionContext(context.TreeContext): def get_yield_values(self): for_parents = [(y, tree.search_ancestor(y, ('for_stmt', 'funcdef', 'while_stmt', 'if_stmt'))) - for y in self.tree_node.yields] + for y in self.tree_node.iter_yield_exprs()] # Calculate if the yields are placed within the same for loop. yields_order = [] diff --git a/jedi/parser/python/tree.py b/jedi/parser/python/tree.py index 564f4fc0..61b355d7 100644 --- a/jedi/parser/python/tree.py +++ b/jedi/parser/python/tree.py @@ -239,9 +239,6 @@ class Scope(PythonBaseNode, DocstringMixin): def __init__(self, children): super(Scope, self).__init__(children) - def iter_return_stmts(self): - return self._search_in_scope(ReturnStmt) - def iter_funcdefs(self): return self._search_in_scope(Function) @@ -461,13 +458,15 @@ class Function(ClassOrFunc): def name(self): return self.children[1] # First token after `def` - @property - def yields(self): + def iter_yield_exprs(self): # TODO This is incorrect, yields are also possible in a statement. - return list(self._search_in_scope(YieldExpr)) + return self._search_in_scope(YieldExpr) + + def iter_return_stmts(self): + return self._search_in_scope(ReturnStmt) def is_generator(self): - return bool(self.yields) + return next(self.iter_yield_exprs(), None) is not None @property def annotation(self): diff --git a/test/test_parser/test_parser_tree.py b/test/test_parser/test_parser_tree.py index 7e4216fc..1451231a 100644 --- a/test/test_parser/test_parser_tree.py +++ b/test/test_parser/test_parser_tree.py @@ -52,12 +52,9 @@ class TestsFunctionAndLambdaParsing(object): assert node.is_generator() is expected.get('is_generator', False) def test_yields(self, node, expected): - # TODO: There's a comment in the code noting that the current implementation is incorrect. This returns an - # empty list at the moment (not e.g. False). - if expected.get('yields', False): - assert node.yields - else: - assert not node.yields + # TODO: There's a comment in the code noting that the current + # implementation is incorrect. + assert node.is_generator() == expected.get('yields', False) def test_annotation(self, node, expected): expected_annotation = expected.get('annotation', None)