Fix crash when analyzer encounters listcomp/genexprs

CHANGELOG.rst

@@ -8,6 +8,8 @@ Next Release (TBD)
 * Fix deployment issue for projects deployed with versions
   prior to 0.10.0
   (`#387 <https://github.com/awslabs/chalice/issues/387>`__)
+* Fix crash in analyzer when encountering genexprs and listcomps
+  (`#263 <https://github.com/awslabs/chalice/issues/263>`__)
 
 
 0.10.0

chalice/analyzer.py

@@ -36,7 +36,7 @@
 import ast
 import symtable
 
-from typing import Dict, Set, Any, Optional, List  # noqa
+from typing import Dict, Set, Any, Optional, List, Union  # noqa
 
 
 APICallT = Dict[str, Set[str]]
@@ -267,6 +267,10 @@ def lookup_sub_namespace(self, name):
                 return self.__class__(child, self._global_table)
         raise ValueError("Unknown symbol name: %s" % name)
 
+    def get_sub_namespaces(self):
+        # type: () -> List[symtable.SymbolTable]
+        return self._local_table.get_children()
+
     def get_name(self):
         # type: () -> str
         return self._local_table.get_name()
@@ -544,6 +548,48 @@ def visit_Return(self, node):
             self._set_inferred_type_for_node(self._current_ast_namespace,
                                              inferred_func_type)
 
+    def visit_ListComp(self, node):
+        # type: (ast.ListComp) -> None
+        # 'listcomp' is the string literal used by python
+        # to creating the SymbolTable for the corresponding
+        # list comp function.
+        self._handle_comprehension(node, 'listcomp')
+
+    def visit_GeneratorExp(self, node):
+        # type: (ast.GeneratorExp) -> None
+        # Generator expressions are an interesting case.
+        # They create a new sub scope, but they're not
+        # explicitly named.  Python just creates a table
+        # with the name "genexpr".
+        self._handle_comprehension(node, 'genexpr')
+
+    def _handle_comprehension(self, node, comprehension_type):
+        # type: (Union[ast.ListComp, ast.GeneratorExp], str) -> None
+        child_scope = self._get_matching_sub_namespace(comprehension_type)
+        if child_scope is None:
+            # If there's no child scope (listcomps in py2) then we can
+            # just analyze the node.elt node in the current scope instead
+            # of creating a new child scope.
+            self.visit(node.elt)
+            return
+        child_table = self._symbol_table.new_sub_table(child_scope)
+        child_infer = self._new_inference_scope(
+            ParsedCode(node.elt, child_table), self._binder)
+        child_infer.bind_types()
+
+    def _get_matching_sub_namespace(self, name):
+        # type: (str) -> symtable.SymbolTable
+        namespaces = [
+            t for t in self._symbol_table.get_sub_namespaces()
+            if t.get_name() == name]
+        if not namespaces:
+            return
+        # We're making a simplification and using the genexpr subnamespace.
+        # This has potential to miss a client call but we don't do
+        # inference on node.generators so this doesn't matter for now.
+        child_scope = namespaces[0]
+        return child_scope
+
     def visit(self, node):
         # type: (Any) -> None
         return ast.NodeVisitor.visit(self, node)

tests/unit/test_analyzer.py

@@ -425,6 +425,39 @@ def test_dict_comp_with_no_client_calls():
     """) == {}
 
 
+def test_can_handle_gen_expr():
+    assert aws_calls("""\
+        import boto3
+        ('a' for y in [1,2,3])
+    """) == {}
+
+
+def test_can_detect_calls_in_gen_expr():
+    assert aws_calls("""\
+        import boto3
+        service_name = 'dynamodb'
+        d = boto3.client('dynamodb')
+        (d.list_tables() for i in [1,2,3])
+    """) == {'dynamodb': set(['list_tables'])}
+
+
+def test_can_detect_calls_in_multiple_gen_exprs():
+    assert aws_calls("""\
+        import boto3
+        d = boto3.client('dynamodb')
+        (d for i in [1,2,3])
+        (d.list_tables() for j in [1,2,3])
+    """) == {'dynamodb': set(['list_tables'])}
+
+
+def test_can_handle_list_expr_with_api_calls():
+    assert aws_calls("""\
+        import boto3
+        d = boto3.client('dynamodb')
+        [d.list_tables() for y in [1,2,3]]
+    """) == {'dynamodb': set(['list_tables'])}
+
+
 #def test_can_handle_dict_comp():
 #    assert aws_calls("""\
 #        import boto3