From 2ba61f305cbac06bb935b00bb593ee8ac63ced42 Mon Sep 17 00:00:00 2001
From: Juha Jeronen <juha.jeronen@jyu.fi>
Date: Thu, 19 Jul 2018 03:04:29 +0300
Subject: [PATCH] fix #3

---
 pyan/analyzer.py | 72 ++++++++++++++++++++++++++++--------------------
 1 file changed, 42 insertions(+), 30 deletions(-)

diff --git a/pyan/analyzer.py b/pyan/analyzer.py
index 0b8baad..0cf9f0d 100644
--- a/pyan/analyzer.py
+++ b/pyan/analyzer.py
@@ -573,28 +573,57 @@ class CallGraphVisitor(ast.NodeVisitor):
 
     def visit_ListComp(self, node):
         self.logger.debug("ListComp")
-        with ExecuteInInnerScope(self, "listcomp"):
-            self.visit(node.elt)
-            self.analyze_generators(node.generators)
+        self.analyze_comprehension(node, "listcomp")
 
     def visit_SetComp(self, node):
         self.logger.debug("SetComp")
-        with ExecuteInInnerScope(self, "setcomp"):
-            self.visit(node.elt)
-            self.analyze_generators(node.generators)
+        self.analyze_comprehension(node, "setcomp")
 
     def visit_DictComp(self, node):
         self.logger.debug("DictComp")
-        with ExecuteInInnerScope(self, "dictcomp"):
-            self.visit(node.key)
-            self.visit(node.value)
-            self.analyze_generators(node.generators)
+        self.analyze_comprehension(node, "dictcomp", field1="key", field2="value")
 
     def visit_GeneratorExp(self, node):
         self.logger.debug("GeneratorExp")
-        with ExecuteInInnerScope(self, "genexpr"):
-            self.visit(node.elt)
-            self.analyze_generators(node.generators)
+        self.analyze_comprehension(node, "genexpr")
+
+    def analyze_comprehension(self, node, label, field1="elt", field2=None):
+        # The outermost iterator is evaluated in the current scope;
+        # everything else in the new inner scope.
+        #
+        # See function symtable_handle_comprehension() in
+        #   https://github.com/python/cpython/blob/master/Python/symtable.c
+        # For how it works, see
+        #   https://stackoverflow.com/questions/48753060/what-are-these-extra-symbols-in-a-comprehensions-symtable
+        # For related discussion, see
+        #   https://bugs.python.org/issue10544
+        gens = node.generators  # tuple of ast.comprehension
+        outermost = gens[0]
+        moregens = gens[1:] if len(gens) > 1 else []
+
+        outermost_iters   = sanitize_exprs(outermost.iter)
+        outermost_targets = sanitize_exprs(outermost.target)
+        for expr in outermost_iters:
+            self.visit(expr)  # set self.last_value (to something and hope for the best)
+
+        with ExecuteInInnerScope(self, label):
+            for expr in outermost_targets:
+                self.visit(expr)  # use self.last_value
+            self.last_value = None
+            for expr in outermost.ifs:
+                self.visit(expr)
+
+            # TODO: there's also an is_async field we might want to use in a future version.
+            for gen in moregens:
+                targets = sanitize_exprs(gen.target)
+                values  = sanitize_exprs(gen.iter)
+                self.analyze_binding(targets, values)
+                for expr in gen.ifs:
+                    self.visit(expr)
+
+            self.visit(getattr(node, field1))  # e.g. node.elt
+            if field2:
+                self.visit(getattr(node, field2))
 
     def visit_Call(self, node):
         self.logger.debug("Call %s" % (get_ast_node_name(node.func)))
@@ -793,23 +822,6 @@ class CallGraphVisitor(ast.NodeVisitor):
                 self.visit(tgt)
             self.last_value = None
 
-    def analyze_generators(self, generators):
-        """Analyze the generators in a comprehension form.
-
-        Analyzes the binding part, and visits the "if" expressions (if any).
-
-        generators: an iterable of ast.comprehension objects
-        """
-
-        for gen in generators:
-            # TODO: there's also an is_async field we might want to use in a future version.
-            targets = sanitize_exprs(gen.target)
-            values  = sanitize_exprs(gen.iter)
-            self.analyze_binding(targets, values)
-
-            for expr in gen.ifs:
-                self.visit(expr)
-
     def resolve_builtins(self, ast_node):
         """Resolve those calls to built-in functions whose return values
         can be determined in a simple manner.