Add iterable checking for all types of comprehensions.

Specifically, list, dict and set comprehensions along with generator
expressions.

2 files changed, 36 insertions, 0 deletions

diff --git a/pylint/checkers/base.py b/pylint/checkers/base.py
index 84f5a8d..7f6f44c 100644
--- a/pylint/checkers/base.py
+++ b/pylint/checkers/base.py
@@ -1533,6 +1533,27 @@ class IterableChecker(_BasicChecker):
         for kwarg in node.kwargs:
             self._check_mapping(kwarg.value, node)
 
+    @check_messages('not-an-iterable')
+    def visit_listcomp(self, node):
+        for gen in node.generators:
+            self._check_iterable(gen.iter, node)
+
+    @check_messages('not-an-iterable')
+    def visit_dictcomp(self, node):
+        for gen in node.generators:
+            self._check_iterable(gen.iter, node)
+
+    @check_messages('not-an-iterable')
+    def visit_setcomp(self, node):
+        for gen in node.generators:
+            self._check_iterable(gen.iter, node)
+
+    @check_messages('not-an-iterable')
+    def visit_generatorexp(self, node):
+        import code; code.interact(local=locals())
+        for gen in node.generators:
+            self._check_iterable(gen.iter, node)
+
 
 def register(linter):
     """required method to auto register this checker"""
diff --git a/pylint/test/unittest_checker_base.py b/pylint/test/unittest_checker_base.py
index fb595e3..b58d8d9 100644
--- a/pylint/test/unittest_checker_base.py
+++ b/pylint/test/unittest_checker_base.py
@@ -398,6 +398,21 @@ class IterableTest(CheckerTestCase):
         with self.assertNoMessages():
             self.checker.visit_call(node)
 
+    def test_non_iterable_in_listcomp(self):
+        node = test_utils.extract_node("""
+        [x for x in 123]
+        """)
+        message = Message('not-an-iterable', node=node, args='123')
+        with self.assertAddsMessages(message):
+            self.checker.visit_listcomp(node)
+        # TODO: more tests
+
+    def test_non_iterable_in_generator(self):
+        node = test_utils.extract_node("__(x for x in 123)")
+        message = Message('not-an-iterable', node=node, args='123')
+        with self.assertAddsMessages(message):
+            self.walk(node.root())
+        # TODO: more tests
 
 
 if __name__ == '__main__':