Merge pull request #11 from attilaolah/master

Catch return with arguments inside generator

3 files changed, 86 insertions, 0 deletions

diff --git a/pyflakes/checker.py b/pyflakes/checker.py
index 1cab12e..f45bb86 100644
--- a/pyflakes/checker.py
+++ b/pyflakes/checker.py
@@ -590,6 +590,38 @@ class Checker(object):
                 self.offset = node_offset
         self.popScope()
 
+    def findReturnWithArgument(self, node):
+        """
+        Finds and returns a return statment that has an argument.
+
+        Note that we should use node.returns in Python 3, but this method is
+        never called in Python 3 so we don't bother checking.
+        """
+        for item in node.body:
+            if isinstance(item, ast.Return) and item.value:
+                return item
+            if hasattr(item, 'body'):
+                found = self.findReturnWithArgument(item)
+                if found is not None:
+                    return found
+
+    def isGenerator(self, node):
+        """
+        Checks whether a function is a generator by looking for a yield
+        statement or expression.
+        """
+        if not isinstance(node.body, list):
+            # lambdas can not be generators
+            return False
+        for item in node.body:
+            if isinstance(item, (ast.Assign, ast.Expr)):
+                if isinstance(item.value, ast.Yield):
+                    return True
+            if hasattr(item, 'body'):
+                if self.isGenerator(item):
+                    return True
+        return False
+
     def ignore(self, node):
         pass
 
@@ -772,6 +804,19 @@ class Checker(object):
                 for name, binding in self.scope.unusedAssignments():
                     self.report(messages.UnusedVariable, binding.source, name)
             self.deferAssignment(checkUnusedAssignments)
+
+            if PY2:
+                def checkReturnWithArgumentInsideGenerator():
+                    """
+                    Check to see if there are any return statements with
+                    arguments but the function is a generator.
+                    """
+                    if self.isGenerator(node):
+                        stmt = self.findReturnWithArgument(node)
+                        if stmt is not None:
+                            self.report(messages.ReturnWithArgsInsideGenerator,
+                                stmt)
+                self.deferAssignment(checkReturnWithArgumentInsideGenerator)
             self.popScope()
 
         self.deferFunction(runFunction)
diff --git a/pyflakes/messages.py b/pyflakes/messages.py
index b7cc177..1f799ec 100644
--- a/pyflakes/messages.py
+++ b/pyflakes/messages.py
@@ -126,3 +126,10 @@ class UnusedVariable(Message):
     def __init__(self, filename, loc, names):
         Message.__init__(self, filename, loc)
         self.message_args = (names,)
+
+
+class ReturnWithArgsInsideGenerator(Message):
+    """
+    Indicates a return statement with arguments inside a generator.
+    """
+    message = '\'return\' with argument inside generator'
diff --git a/pyflakes/test/test_return_with_arguments_inside_generator.py b/pyflakes/test/test_return_with_arguments_inside_generator.py
new file mode 100644
index 0000000..8ec8234
--- /dev/null
+++ b/pyflakes/test/test_return_with_arguments_inside_generator.py
@@ -0,0 +1,34 @@
+
+from sys import version_info
+
+from pyflakes import messages as m
+from pyflakes.test.harness import TestCase, skipIf
+
+
+class Test(TestCase):
+    @skipIf(version_info >= (3,), 'new in Python 3')
+    def test_return(self):
+        self.flakes('''
+        class a:
+            def b():
+                for x in a.c:
+                    if x:
+                        yield x
+                return a
+        ''', m.ReturnWithArgsInsideGenerator)
+
+    @skipIf(version_info >= (3,), 'new in Python 3')
+    def test_returnNone(self):
+        self.flakes('''
+        def a():
+            yield 12
+            return None
+        ''', m.ReturnWithArgsInsideGenerator)
+
+    @skipIf(version_info >= (3,), 'new in Python 3')
+    def test_returnYieldExpression(self):
+        self.flakes('''
+        def a():
+            b = yield a
+            return b
+        ''', m.ReturnWithArgsInsideGenerator)