Fix "continue" and "break" checks ignoring py3.5's "async for" loop (#71)

2 files changed, 64 insertions, 1 deletions

diff --git a/pyflakes/checker.py b/pyflakes/checker.py
index 43acc69..9f354fd 100644
--- a/pyflakes/checker.py
+++ b/pyflakes/checker.py
@@ -12,6 +12,7 @@ import sys
 PY2 = sys.version_info < (3, 0)
 PY32 = sys.version_info < (3, 3)    # Python 2.5 to 3.2
 PY33 = sys.version_info < (3, 4)    # Python 2.5 to 3.3
+PY34 = sys.version_info < (3, 5)    # Python 2.5 to 3.4
 try:
     sys.pypy_version_info
     PYPY = True
@@ -55,6 +56,11 @@ else:
         if isinstance(n, ast.Try):
             return [n.body + n.orelse] + [[hdl] for hdl in n.handlers]
 
+if PY34:
+    LOOP_TYPES = (ast.While, ast.For)
+else:
+    LOOP_TYPES = (ast.While, ast.For, ast.AsyncFor)
+
 
 class _FieldsOrder(dict):
     """Fix order of AST node fields."""
@@ -943,7 +949,7 @@ class Checker(object):
         n = node
         while hasattr(n, 'parent'):
             n, n_child = n.parent, n
-            if isinstance(n, (ast.While, ast.For)):
+            if isinstance(n, LOOP_TYPES):
                 # Doesn't apply unless it's in the loop itself
                 if n_child not in n.orelse:
                     return
diff --git a/pyflakes/test/test_other.py b/pyflakes/test/test_other.py
index ae6cea2..7efc1c5 100644
--- a/pyflakes/test/test_other.py
+++ b/pyflakes/test/test_other.py
@@ -1720,6 +1720,63 @@ class TestAsyncStatements(TestCase):
         ''')
 
     @skipIf(version_info < (3, 5), 'new in Python 3.5')
+    def test_loopControlInAsyncFor(self):
+        self.flakes('''
+        async def read_data(db):
+            output = []
+            async for row in db.cursor():
+                if row[0] == 'skip':
+                    continue
+                output.append(row)
+            return output
+        ''')
+
+        self.flakes('''
+        async def read_data(db):
+            output = []
+            async for row in db.cursor():
+                if row[0] == 'stop':
+                    break
+                output.append(row)
+            return output
+        ''')
+
+    @skipIf(version_info < (3, 5), 'new in Python 3.5')
+    def test_loopControlInAsyncForElse(self):
+        self.flakes('''
+        async def read_data(db):
+            output = []
+            async for row in db.cursor():
+                output.append(row)
+            else:
+                continue
+            return output
+        ''', m.ContinueOutsideLoop)
+
+        self.flakes('''
+        async def read_data(db):
+            output = []
+            async for row in db.cursor():
+                output.append(row)
+            else:
+                break
+            return output
+        ''', m.BreakOutsideLoop)
+
+    @skipIf(version_info < (3, 5), 'new in Python 3.5')
+    def test_continueInAsyncForFinally(self):
+        self.flakes('''
+        async def read_data(db):
+            output = []
+            async for row in db.cursor():
+                try:
+                    output.append(row)
+                finally:
+                    continue
+            return output
+        ''', m.ContinueInFinally)
+
+    @skipIf(version_info < (3, 5), 'new in Python 3.5')
     def test_asyncWith(self):
         self.flakes('''
         async def commit(session, data):