Add more positive functional tests for the iterable context checker

2 files changed, 79 insertions, 25 deletions

diff --git a/pylint/test/functional/iterable_context.py b/pylint/test/functional/iterable_context.py
index 8282683..0571583 100644
--- a/pylint/test/functional/iterable_context.py
+++ b/pylint/test/functional/iterable_context.py
@@ -2,7 +2,7 @@
 Checks that primitive values are not used in an
 iterating/mapping context.
 """
-# pylint: disable=missing-docstring,invalid-name
+# pylint: disable=missing-docstring,invalid-name,too-few-public-methods,no-init
 from __future__ import print_function
 
 # for-statement
@@ -12,25 +12,87 @@ for i in 42:  # [not-an-iterable]
 for i in True:  # [not-an-iterable]
     pass
 
-# funcall-starargs
-def test(*args, **kwargs):
-    print(args, kwargs)
+for i in range(10):
+    pass
+
+for i in ''.join(x ** 2 for x in range(10)):
+    pass
+
+numbers = [1, 2, 3]
+inumbers = iter(numbers)
+
+for i in inumbers:
+    pass
+
+for i in "123":
+    pass
+
+for i in u"123":
+    pass
+
+for i in set(numbers):
+    pass
+
+for i in frozenset(numbers):
+    pass
+
+for i in dict(a=1, b=2):
+    pass
+
+for i in [x for x in range(10)]:
+    pass
 
-test(*1)  # [not-an-iterable]
-test(*False)  # [not-an-iterable]
+for i in {x for x in range(1, 100, 2)}:
+    pass
+
+for i in {x: 10 - x for x in range(10)}:
+    pass
+
+def powers_of_two():
+    k = 0
+    while k < 10:
+        yield 2 ** k
+        k += 1
+
+for i in powers_of_two():
+    pass
+
+
+# check for old-style iterator
+class C(object):
+    def __getitem__(self, k):
+        if k > 10:
+            raise IndexError
+        return k + 1
+
+    def __len__(self):
+        return 10
 
-# funcall-kwargs
-test(**1)  # [not-a-mapping]
-test(**None)  # [not-a-mapping]
+for i in C():
+    print(i)
 
-# list-comprehension
-test([3 ** x for x in 10])  # [not-an-iterable]
 
-# dict-comprehension
-test({k: chr(k) for k in 128})  # [not-an-iterable]
+# check for custom iterators
+class A:
+    pass
+
+
+class B:
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        return 1
+
+    def next(self):
+        return 1
 
-# set-comprehension
-test({x for x in 32})  # [not-an-iterable]
 
-# generator-expression
-test(str(x) for x in 10)  # [not-an-iterable]
+def test(*args):
+    pass
+
+
+test(*A())  # [not-an-iterable]
+test(*B())
+for i in A():  # [not-an-iterable]
+    pass
diff --git a/pylint/test/functional/iterable_context.txt b/pylint/test/functional/iterable_context.txt
index 8c9f098..b618ffb 100644
--- a/pylint/test/functional/iterable_context.txt
+++ b/pylint/test/functional/iterable_context.txt
@@ -1,10 +1,2 @@
 not-an-iterable:9::Non-iterable value 42 is used in an iterating context
 not-an-iterable:12::Non-iterable value True is used in an iterating context
-not-an-iterable:19::Non-iterable value 1 is used in an iterating context
-not-an-iterable:20::Non-iterable value False is used in an iterating context
-not-a-mapping:23::Non-mapping value 1 is used in a mapping context
-not-a-mapping:24::Non-mapping value None is used in a mapping context
-not-an-iterable:27::Non-iterable value 10 is used in an iterating context
-not-an-iterable:30::Non-iterable value 128 is used in an iterating context
-not-an-iterable:33::Non-iterable value 32 is used in an iterating context
-not-an-iterable:36::Non-iterable value 10 is used in an iterating context