Add check for non-iterator returned by __iter__.

3 files changed, 92 insertions, 1 deletions

diff --git a/checkers/classes.py b/checkers/classes.py
index 98a2f27..f8765a2 100644
--- a/checkers/classes.py
+++ b/checkers/classes.py
@@ -16,15 +16,18 @@
 """classes checker for Python code
 """
 from __future__ import generators
-
+import sys
 import astroid
 from astroid import YES, Instance, are_exclusive, AssAttr
+from astroid.bases import Generator
 
 from pylint.interfaces import IAstroidChecker
 from pylint.checkers import BaseChecker
 from pylint.checkers.utils import (PYMETHODS, overrides_a_method,
     check_messages, is_attr_private, is_attr_protected, node_frame_class)
 
+_PY3K = sys.version_info >= (3, 0)
+
 def class_is_abstract(node):
     """return true if the given class node should be considered as an abstract
     class
@@ -142,6 +145,11 @@ MSGS = {
               'non-parent-init-called',
               'Used when an __init__ method is called on a class which is not \
               in the direct ancestors for the analysed class.'),
+    'W0234': ('__iter__ returns non-iterator',
+              'non-iterator-returned',
+              'Used when an __iter__ method returns something which is not an \
+               iterator.'),
+
 
     }
 
@@ -311,6 +319,34 @@ a metaclass class method.'}
         except astroid.NotFoundError:
             pass
 
+        # check non-iterators in __iter__
+        if node.name == '__iter__':
+            self._check_iter(node)
+
+    def _check_iter(self, node):
+        try:
+            infered = node.infer_call_result(node)
+        except astroid.InferenceError:
+            return
+
+        if _PY3K:
+            next = '__next__'
+        else:
+            next = 'next'
+        for infered_node in infered:
+            if (infered_node is YES
+                or isinstance(infered_node, Generator)):
+                continue
+            if isinstance(infered_node, astroid.Instance):
+                for meth in ('__iter__', next):
+                    try:
+                        infered_node.local_attr(meth)
+                    except astroid.NotFoundError:
+                        self.add_message('non-iterator-returned',
+                                         node=node)
+                        break
+
+
     def leave_function(self, node):
         """on method node, check if this method couldn't be a function
 
diff --git a/test/input/func_non_iterator_returned.py b/test/input/func_non_iterator_returned.py
new file mode 100644
index 0000000..85a61c2
--- /dev/null
+++ b/test/input/func_non_iterator_returned.py
@@ -0,0 +1,52 @@
+"""Check non-iterators returned by __iter__ """
+
+# pylint: disable-msg=too-few-public-methods
+
+__revision__ = 0
+
+class FirstGoodIterator(object):
+    """ yields in iterator. """
+
+    def __iter__(self):
+        for index in range(10):
+            yield index
+
+class SecondGoodIterator(object):
+    """ __iter__ and next """
+
+    def __iter__(self):
+        return self
+
+    def next(self): # pylint: disable-msg=no-self-use
+        """ Infinite iterator, but still an iterator """
+        return 1
+
+class ThirdGoodIterator(object):
+    """ Returns other iterator, not the current instance """
+
+    def __iter__(self):
+        return SecondGoodIterator()
+
+class FourthGoodIterator(object):
+    """ __iter__ returns iter(...) """
+
+    def __iter__(self):
+        return iter(range(10))
+
+class FirstBadIterator(object):
+    """ __iter__ returns a list """
+
+    def __iter__(self):
+        return []
+
+class SecondBadIterator(object):
+    """ __iter__ without next """
+
+    def __iter__(self):
+        return self
+
+class ThirdBadIterator(object):
+    """ __iter__ returns an instance of another non-iterator """
+
+    def __iter__(self):
+        return SecondBadIterator()
diff --git a/test/messages/func_non_iterator_returned.txt b/test/messages/func_non_iterator_returned.txt
new file mode 100644
index 0000000..821527c
--- /dev/null
+++ b/test/messages/func_non_iterator_returned.txt
@@ -0,0 +1,3 @@
+W: 39:FirstBadIterator.__iter__: __iter__ returns non-iterator

+W: 45:SecondBadIterator.__iter__: __iter__ returns non-iterator

+W: 51:ThirdBadIterator.__iter__: __iter__ returns non-iterator