Improve iterable inference and add more unit tests for iterable checker

2 files changed, 62 insertions, 11 deletions

diff --git a/pylint/checkers/base.py b/pylint/checkers/base.py
index 7f6f44c..904cf9f 100644
--- a/pylint/checkers/base.py
+++ b/pylint/checkers/base.py
@@ -1499,15 +1499,25 @@ class IterableChecker(_BasicChecker):
         is_string = self._is_string(value)
         return is_iterable_value or is_string
 
+    def _not_an_iterable(self, value):
+        is_iterable = self._is_iterable(value)
+        is_const = isinstance(value, astroid.Const)
+        return not is_iterable and is_const
+
     def _is_mapping(self, value):
-        return isinstance(value, astroid.Dict)
+        return isinstance(value, (astroid.Dict, astroid.DictComp))
+
+    def _not_a_mapping(self, value):
+        is_mapping = self._is_mapping(value)
+        is_const = isinstance(value, astroid.Const)
+        return not is_mapping and is_const
 
     def _check_iterable(self, element, root_node):
         infered = helpers.safe_infer(element)
         if infered is not None:
-            if not self._is_iterable(infered):
+            if self._not_an_iterable(infered):
                 self.add_message('not-an-iterable',
-                                 args=infered.as_string(),
+                                 args=element.as_string(),
                                  node=root_node)
 
     def _check_mapping(self, element, root_node):
@@ -1515,7 +1525,7 @@ class IterableChecker(_BasicChecker):
         if infered is not None:
             if not self._is_mapping(infered):
                 self.add_message('not-a-mapping',
-                                 args=infered.as_string(),
+                                 args=element.as_string(),
                                  node=root_node)
 
     @check_messages('not-an-iterable')
@@ -1550,7 +1560,6 @@ class IterableChecker(_BasicChecker):
 
     @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)
 
diff --git a/pylint/test/unittest_checker_base.py b/pylint/test/unittest_checker_base.py
index b58d8d9..538cf94 100644
--- a/pylint/test/unittest_checker_base.py
+++ b/pylint/test/unittest_checker_base.py
@@ -369,7 +369,7 @@ class IterableTest(CheckerTestCase):
         with self.assertNoMessages():
             self.walk(node.root())
 
-    def test_non_mapping_in_kwargs(self):
+    def test_non_mapping_in_funcall_kwargs(self):
         node = test_utils.extract_node("""
         foo(**123)
         """)
@@ -381,7 +381,7 @@ class IterableTest(CheckerTestCase):
         args = [1, 2, 3]
         foo(**args)
         """)
-        message = Message('not-a-mapping', node=node, args='[1, 2, 3]')
+        message = Message('not-a-mapping', node=node, args='args')
         with self.assertAddsMessages(message):
             self.walk(node.root())
 
@@ -398,21 +398,63 @@ class IterableTest(CheckerTestCase):
         with self.assertNoMessages():
             self.checker.visit_call(node)
 
+        node = test_utils.extract_node("""
+        foo(**{str(x): x for x in range(5)})
+        """)
+        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]
+        [x ** 2 for x in 10]
         """)
-        message = Message('not-an-iterable', node=node, args='123')
+        message = Message('not-an-iterable', node=node, args='10')
         with self.assertAddsMessages(message):
             self.checker.visit_listcomp(node)
-        # TODO: more tests
+
+        node = test_utils.extract_node("""
+        [x ** 2 for x in range(10)]
+        """)
+        with self.assertNoMessages():
+            self.checker.visit_listcomp(node)
+
+    def test_non_iterable_in_dictcomp(self):
+        node = test_utils.extract_node("""
+        {x: chr(x) for x in 256}
+        """)
+        message = Message('not-an-iterable', node=node, args='256')
+        with self.assertAddsMessages(message):
+            self.checker.visit_dictcomp(node)
+
+        node = test_utils.extract_node("""
+        {ord(x): x for x in "aoeui"}
+        """)
+        with self.assertNoMessages():
+            self.checker.visit_dictcomp(node)
+
+    def test_non_iterable_in_setcomp(self):
+        node = test_utils.extract_node("""
+        {2 ** x for x in 10}
+        """)
+        message = Message('not-an-iterable', node=node, args='10')
+        with self.assertAddsMessages(message):
+            self.checker.visit_setcomp(node)
+
+        node = test_utils.extract_node("""
+        {2 ** x for x in range(10)}
+        """)
+        with self.assertNoMessages():
+            self.checker.visit_setcomp(node)
 
     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
+
+        node = test_utils.extract_node("__(chr(x) for x in range(25))")
+        with self.assertNoMessages():
+            self.walk(node.root())
 
 
 if __name__ == '__main__':