Repository: PyCQA/pyflakes
Pull request #339 (2018-06-27 14:32:37)

Handle subtrees as input to checker
Closes https://github.com/PyCQA/pyflakes/issues/338

MacBox7: This fixes the issue with minimum changes. The problem i.e being fixed here is in regard to the node to which the `Argument` node is being associated.

Previously the `source` attribute of `Argument` node i.e `node.source` pointed to `FunctionDef` node resulting in incorrect `lineno` and `column_offset` of the argument. But now the source attribute points to the `arg` node thus associating arguments with correct lineno and col_offset. This fix is only for PY3 as `arg` node was introduced in here. Patch has no changes for PY2.
MacBox7: Just added some tests to make sure `ARG` and `ARGUMENT` handler are handled. With all possible permutations of the valid ways by which arguments can be supplied to the function.
MacBox7: New changes done to do PR:
1. Initially changes were only done for python3 arguments. Now more changes have been done to fix the arguments of python2 as well. The location information(line, col, node) in python2 now corresponds to ast.Name object for each respective argument. Initially, they were all pointing to `ast.FunctionDef` which was wrong and led to incorrect positional information. However `kwarg` and `vararg` still point to `ast.Function` as they are represented as plain string in python2.

2. Added a test harness that can be used to test code snippets rather than complete code.
MacBox7: Need opinion on whether representing `kwarg` and `vararg` with `ast.argument` is a good choice.

Patch:
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


Test patch:
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):
     """