Handle subtrees as input to checker

This patch essentially adds the the
functionality for the Checker to handle
subtrees. Along with this args have also
been parsed through handlers.

Closes #338

Author: ankitxjoshi

pyflakes/checker.py

@@ -95,9 +95,15 @@ def iter_child_nodes(node, omit=None, _fields_order=_FieldsOrder()):
     """
     Yield all direct child nodes of *node*, that is, all fields that
     are nodes and all items of fields that are lists of nodes.
+
+    :param node:          AST node to be iterated upon
+    :param omit:          String or tuple of strings denoting the
+                          attributes of the node to be omitted from
+                          further parsing
+    :param _fields_order: Order of AST node fields
     """
     for name in _fields_order[node.__class__]:
-        if name == omit:
+        if omit and name in omit:
             continue
         field = getattr(node, name, None)
         if isinstance(field, ast.AST):
@@ -472,6 +478,17 @@ class Checker(object):
         callables which are deferred assignment checks.
     """
 
+    _ast_node_scope = {
+        ast.Module: ModuleScope,
+        ast.ClassDef: ClassScope,
+        ast.FunctionDef: FunctionScope,
+        ast.Lambda: FunctionScope,
+        ast.ListComp: GeneratorScope,
+        ast.SetComp: GeneratorScope,
+        ast.GeneratorExp: GeneratorScope,
+        ast.DictComp: GeneratorScope,
+    }
+
     nodeDepth = 0
     offset = None
     traceTree = False
@@ -493,7 +510,10 @@ def __init__(self, tree, filename='(none)', builtins=None,
         if builtins:
             self.builtIns = self.builtIns.union(builtins)
         self.withDoctest = withDoctest
-        self.scopeStack = [ModuleScope()]
+        try:
+            self.scopeStack = [Checker._ast_node_scope[type(tree)]()]
+        except KeyError:
+            raise RuntimeError('No scope implemented for the node %r' % tree)
         self.exceptHandlers = [()]
         self.root = tree
         self.handleChildren(tree)
@@ -643,6 +663,18 @@ def descendantOf(self, node, ancestors, stop):
                 return True
         return False
 
+    def _getAncestor(self, node, ancestor_type):
+        parent = node
+        while True:
+            if parent is self.root:
+                return None
+            parent = self.getParent(parent)
+            if isinstance(parent, ancestor_type):
+                return parent
+
+    def getScopeNode(self, node):
+        return self._getAncestor(node, tuple(Checker._ast_node_scope.keys()))
+
     def differentForks(self, lnode, rnode):
         """True, if lnode and rnode are located on different forks of IF/TRY"""
         ancestor = self.getCommonAncestor(lnode, rnode, self.root)
@@ -790,6 +822,8 @@ def handleNodeStore(self, node):
             binding = Binding(name, node)
         elif name == '__all__' and isinstance(self.scope, ModuleScope):
             binding = ExportBinding(name, node.parent, self.scope)
+        elif isinstance(getattr(node, 'ctx', None), ast.Param):
+            binding = Argument(name, self.getScopeNode(node))
         else:
             binding = Assignment(name, node)
         self.addBinding(node, binding)
@@ -1103,13 +1137,12 @@ def NAME(self, node):
                     and isinstance(node.parent, ast.Call)):
                 # we are doing locals() call in current scope
                 self.scope.usesLocals = True
-        elif isinstance(node.ctx, (ast.Store, ast.AugStore)):
+        elif isinstance(node.ctx, (ast.Store, ast.AugStore, ast.Param)):
             self.handleNodeStore(node)
         elif isinstance(node.ctx, ast.Del):
             self.handleNodeDelete(node)
         else:
-            # must be a Param context -- this only happens for names in function
-            # arguments, but these aren't dispatched through here
+            # Unknown context
             raise RuntimeError("Got impossible expression context: %r" % (node.ctx,))
 
     def CONTINUE(self, node):
@@ -1225,15 +1258,8 @@ def addArgs(arglist):
         def runFunction():
 
             self.pushScope()
-            for name in args:
-                self.addBinding(node, Argument(name, node))
-            if isinstance(node.body, list):
-                # case for FunctionDefs
-                for stmt in node.body:
-                    self.handleNode(stmt, node)
-            else:
-                # case for Lambdas
-                self.handleNode(node.body, node)
+
+            self.handleChildren(node, omit='decorator_list')
 
             def checkUnusedAssignments():
                 """
@@ -1257,6 +1283,18 @@ def checkReturnWithArgumentInsideGenerator():
 
         self.deferFunction(runFunction)
 
+    def ARGUMENTS(self, node):
+        self.handleChildren(node, omit=('defaults', 'kw_defaults'))
+        if PY2:
+            scope_node = self.getScopeNode(node)
+            if node.vararg:
+                self.addBinding(node, Argument(node.vararg, scope_node))
+            if node.kwarg:
+                self.addBinding(node, Argument(node.kwarg, scope_node))
+
+    def ARG(self, node):
+        self.addBinding(node, Argument(node.arg, self.getScopeNode(node)))
+
     def CLASSDEF(self, node):
         """
         Check names used in a class definition, including its decorators, base

pyflakes/test/harness.py

@@ -17,6 +17,9 @@ class TestCase(unittest.TestCase):
 
     def flakes(self, input, *expectedOutputs, **kw):
         tree = compile(textwrap.dedent(input), "<test>", "exec", PyCF_ONLY_AST)
+        if kw.get('is_segment'):
+            tree = tree.body[0]
+            kw.pop('is_segment')
         w = checker.Checker(tree, withDoctest=self.withDoctest, **kw)
         outputs = [type(o) for o in w.messages]
         expectedOutputs = list(expectedOutputs)

pyflakes/test/test_code_segment.py

@@ -0,0 +1,126 @@
+from pyflakes import messages as m
+from pyflakes.checker import (FunctionScope, ClassScope, ModuleScope,
+                              Argument, FunctionDefinition, Assignment)
+from pyflakes.test.harness import TestCase
+
+
+class TestCodeSegments(TestCase):
+    """
+    Tests for segments of a module
+    """
+
+    def test_function_segment(self):
+        self.flakes('''
+        def foo():
+            def bar():
+                pass
+        ''', is_segment=True)
+
+        self.flakes('''
+        def foo():
+            def bar():
+                x = 0
+        ''', m.UnusedVariable, is_segment=True)
+
+    def test_class_segment(self):
+        self.flakes('''
+        class Foo:
+            class Bar:
+                pass
+        ''', is_segment=True)
+
+        self.flakes('''
+        class Foo:
+            def bar():
+                x = 0
+        ''', m.UnusedVariable, is_segment=True)
+
+    def test_scope_class(self):
+        checker = self.flakes('''
+        class Foo:
+            x = 0
+            def bar(a, b=1, *d, **e):
+                pass
+        ''', is_segment=True)
+
+        scopes = checker.deadScopes
+        module_scopes = [
+            scope for scope in scopes if scope.__class__ is ModuleScope]
+        class_scopes = [
+            scope for scope in scopes if scope.__class__ is ClassScope]
+        function_scopes = [
+            scope for scope in scopes if scope.__class__ is FunctionScope]
+
+        # Ensure module scope is not present because we are analysing
+        # the inner contents of Foo
+        self.assertEqual(len(module_scopes), 0)
+        self.assertEqual(len(class_scopes), 1)
+        self.assertEqual(len(function_scopes), 1)
+
+        class_scope = class_scopes[0]
+        function_scope = function_scopes[0]
+
+        self.assertIsInstance(class_scope, ClassScope)
+        self.assertIsInstance(function_scope, FunctionScope)
+
+        self.assertIn('x', class_scope)
+        self.assertIn('bar', class_scope)
+
+        self.assertIn('a', function_scope)
+        self.assertIn('b', function_scope)
+        self.assertIn('d', function_scope)
+        self.assertIn('e', function_scope)
+
+        self.assertIsInstance(class_scope['bar'], FunctionDefinition)
+        self.assertIsInstance(class_scope['x'], Assignment)
+
+        self.assertIsInstance(function_scope['a'], Argument)
+        self.assertIsInstance(function_scope['b'], Argument)
+        self.assertIsInstance(function_scope['d'], Argument)
+        self.assertIsInstance(function_scope['e'], Argument)
+
+    def test_scope_function(self):
+        checker = self.flakes('''
+        def foo(a, b=1, *d, **e):
+            def bar(f, g=1, *h, **i):
+                pass
+        ''', is_segment=True)
+
+        scopes = checker.deadScopes
+        module_scopes = [
+            scope for scope in scopes if scope.__class__ is ModuleScope]
+        function_scopes = [
+            scope for scope in scopes if scope.__class__ is FunctionScope]
+
+        # Ensure module scope is not present because we are analysing
+        # the inner contents of foo
+        self.assertEqual(len(module_scopes), 0)
+        self.assertEqual(len(function_scopes), 2)
+
+        function_scope_foo = function_scopes[1]
+        function_scope_bar = function_scopes[0]
+
+        self.assertIsInstance(function_scope_foo, FunctionScope)
+        self.assertIsInstance(function_scope_bar, FunctionScope)
+
+        self.assertIn('a', function_scope_foo)
+        self.assertIn('b', function_scope_foo)
+        self.assertIn('d', function_scope_foo)
+        self.assertIn('e', function_scope_foo)
+        self.assertIn('bar', function_scope_foo)
+
+        self.assertIn('f', function_scope_bar)
+        self.assertIn('g', function_scope_bar)
+        self.assertIn('h', function_scope_bar)
+        self.assertIn('i', function_scope_bar)
+
+        self.assertIsInstance(function_scope_foo['bar'], FunctionDefinition)
+        self.assertIsInstance(function_scope_foo['a'], Argument)
+        self.assertIsInstance(function_scope_foo['b'], Argument)
+        self.assertIsInstance(function_scope_foo['d'], Argument)
+        self.assertIsInstance(function_scope_foo['e'], Argument)
+
+        self.assertIsInstance(function_scope_bar['f'], Argument)
+        self.assertIsInstance(function_scope_bar['g'], Argument)
+        self.assertIsInstance(function_scope_bar['h'], Argument)
+        self.assertIsInstance(function_scope_bar['i'], Argument)

pyflakes/test/test_other.py

@@ -1159,6 +1159,37 @@ def f(): global foo
         def g(): foo = 'anything'; foo.is_used()
         ''')
 
+    def test_function_arguments(self):
+        """
+        Test to traverse ARG and ARGUMENT handler
+        """
+        self.flakes('''
+        def foo(a, b):
+            pass
+        ''')
+
+        self.flakes('''
+        def foo(a, b, c=0):
+            pass
+        ''')
+
+        self.flakes('''
+        def foo(a, b, c=0, *args):
+            pass
+        ''')
+
+        self.flakes('''
+        def foo(a, b, c=0, *args, **kwargs):
+            pass
+        ''')
+
+    @skipIf(version_info < (3, 3), "Python >= 3.3 only")
+    def test_function_arguments_python3(self):
+        self.flakes('''
+        def foo(a, b, c=0, *args, d=0, **kwargs):
+            pass
+        ''')
+
 
 class TestUnusedAssignment(TestCase):
     """