Skip to content

Commit

Permalink
FunctionDef.is_generator properly handles yield nodes in If tests
Browse files Browse the repository at this point in the history
  • Loading branch information
PCManticore committed Jun 9, 2020
1 parent ca821aa commit 09ee106
Show file tree
Hide file tree
Showing 3 changed files with 22 additions and 0 deletions.
4 changes: 4 additions & 0 deletions ChangeLog
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,10 @@ Release Date: TBA

Close PyCQA/pylint#3639

* `FunctionDef.is_generator` properly handles `yield` nodes in `If` tests

Close PyCQA/pylint#3583


What's New in astroid 2.4.2?
============================
Expand Down
5 changes: 5 additions & 0 deletions astroid/node_classes.py
Original file line number Diff line number Diff line change
Expand Up @@ -3508,6 +3508,11 @@ def get_children(self):
def has_elif_block(self):
return len(self.orelse) == 1 and isinstance(self.orelse[0], If)

def _get_yield_nodes_skip_lambdas(self):
"""An If node can contain a Yield node in the test"""
yield from self.test._get_yield_nodes_skip_lambdas()
yield from super()._get_yield_nodes_skip_lambdas()


class IfExp(NodeNG):
"""Class representing an :class:`ast.IfExp` node.
Expand Down
13 changes: 13 additions & 0 deletions tests/unittest_nodes.py
Original file line number Diff line number Diff line change
Expand Up @@ -1347,5 +1347,18 @@ def paused_iter(iterable):
assert bool(node.is_generator())


def test_is_generator_for_yield_in_if():
code = """
import asyncio
def paused_iter(iterable):
if (yield from asyncio.sleep(0.01)):
pass
return
"""
node = astroid.extract_node(code)
assert bool(node.is_generator())


if __name__ == "__main__":
unittest.main()

0 comments on commit 09ee106

Please sign in to comment.