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 https://github.com/PyCQA/pyflakes/issues/338

4 files changed, 212 insertions, 14 deletions

diff --git a/pyflakes/checker.py b/pyflakes/checker.py
index a57ce50..1c30dda 100644
--- a/pyflakes/checker.py
+++ b/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 @@ class Checker(object):
         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 @@ class Checker(object):
                 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 @@ class Checker(object):
             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 @@ class Checker(object):
                     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 @@ class Checker(object):
         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 @@ class Checker(object):
 
         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
diff --git a/pyflakes/test/harness.py b/pyflakes/test/harness.py
index 3f592db..0a58bd5 100644
--- a/pyflakes/test/harness.py
+++ b/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)
diff --git a/pyflakes/test/test_code_segment.py b/pyflakes/test/test_code_segment.py
new file mode 100644
index 0000000..5e53f8f
--- /dev/null
+++ b/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)
diff --git a/pyflakes/test/test_other.py b/pyflakes/test/test_other.py
index 4dbaf23..fa093ce 100644
--- a/pyflakes/test/test_other.py
+++ b/pyflakes/test/test_other.py
@@ -1159,6 +1159,37 @@ class Test(TestCase):
         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):
     """