Merge pull request #2146 from RasmusWL/python-improve-iter-returns-non-iterator

Python: improve py/iter-returns-non-iterator

python/ql/src/Functions/IterReturnsNonIterator.py

@@ -1,18 +1,33 @@
 class MyRange(object):
     def __init__(self, low, high):
-        self.current = low
+        self.low = low
         self.high = high
 
     def __iter__(self):
-        return self
+        return MyRangeIterator(self.low, self.high)
 
-#Fixed version
-class MyRange(object):
-    def __init__(self, low, high):
+    def skip(self, to_skip):
+        return MyRangeIterator(self.low, self.high, to_skip)
+
+class MyRangeIterator(object):
+    def __init__(self, low, high, skip=None):
         self.current = low
         self.high = high
+        self.skip = skip
 
-    def __iter__(self):
-        while self.current < self.high:
-            yield self.current
-            self.current += 1
+    def __next__(self):
+        if self.current >= self.high:
+            raise StopIteration
+        to_return = self.current
+        self.current += 1
+        if self.skip and to_return in self.skip:
+            return self.__next__()
+        return to_return
+
+    # Problem is fixed by uncommenting these lines
+    # def __iter__(self):
+    #     return self
+
+my_range = MyRange(0,10)
+x = sum(my_range) # x = 45
+y = sum(my_range.skip({6,9})) # TypeError: 'MyRangeIterator' object is not iterable

python/ql/src/Functions/IterReturnsNonIterator.qhelp

@@ -3,35 +3,33 @@
   "qhelp.dtd">
 <qhelp>
 <overview>
-<p>The <code>__iter__</code> method of a class should return an iterator.
+<p>The <code>__iter__</code> method of a class should always return an iterator.</p>
 
-Iteration in Python relies on this behavior and attempting to iterate over an
-instance of a class with an incorrect <code>__iter__</code> method will raise a TypeError.
+<p>Iterators must implement both <code>__next__</code> and <code>__iter__</code> for Python 3, or both <code>next</code> and <code>__iter__</code> for Python 2. The <code>__iter__</code> method of the iterator must return the iterator object itself.</p>
+
+<p>Iteration in Python relies on this behavior and attempting to iterate over an instance of a class with an incorrect <code>__iter__</code> method can raise a <code>TypeError</code>.
 </p>
 
-
 </overview>
 <recommendation>
-<p>Make the <code>__iter__</code> return a new iterator, either as an instance of
-a separate class or as a generator.</p>
-
+<p>Make sure the value returned by <code>__iter__</code> implements the full iterator protocol.</p>
 </recommendation>
 <example>
-<p>In this example the <code>MyRange</code> class's <code>__iter__</code> method does not
-return  an iterator. This will cause the program to fail when anyone attempts
-to use the iterator in a <code>for</code> loop or <code>in</code> statement.
-</p>
+<p>In this example, we have implemented our own version of <code>range</code>, extending the  normal functionality with the ability to skip some elements by using the <code>skip</code> method. However, the iterator <code>MyRangeIterator</code> does not fully implement the iterator protocol (namely it is missing <code>__iter__</code>).</p>
 
-<p>The fixed version implements the <code>__iter__</code> method as a generator function.</p>
+<p>Iterating over the elements in the range seems to work on the surface, for example the code <code>x = sum(my_range)</code> gives the expected result. However, if we run <code>sum(iter(my_range))</code> we get a <code>TypeError: 'MyRangeIterator' object is not iterable</code>.</p>
+
+<p>If we try to skip some elements using our custom method, for example <code>y = sum(my_range.skip({6,9}))</code>, this also raises a <code>TypeError</code>.</p>
+
+<p>The fix is to implement the <code>__iter__</code> method in <code>MyRangeIterator</code>.</p>
 
 <sample src="IterReturnsNonIterator.py" />
 
 </example>
 <references>
 
-  <li>Python Language Reference: <a href="http://docs.python.org/2.7/reference/datamodel.html#object.__iter__">object.__iter__</a>.</li>
-  <li>Python Standard Library: <a href="http://docs.python.org/2/library/stdtypes.html#typeiter">Iterator Types</a>.</li>
-
+  <li>Python Language Reference: <a href="https://docs.python.org/3/reference/datamodel.html#object.__iter__">object.__iter__</a>.</li>
+  <li>Python Standard Library: <a href="https://docs.python.org/3/library/stdtypes.html#typeiter">Iterator Types</a>.</li>
 
 </references>
 </qhelp>

python/ql/src/Functions/IterReturnsNonIterator.ql

@@ -12,21 +12,20 @@
 
 import python
 
-FunctionObject iter_method(ClassObject t) {
-		result = t.lookupAttribute("__iter__")
-}
-
-cached ClassObject return_type(FunctionObject f) {
+ClassObject return_type(FunctionObject f) {
     exists(ControlFlowNode n, Return ret |
-        ret.getScope() = f.getFunction() and ret.getValue() = n.getNode() and
+        ret.getScope() = f.getFunction() and
+        ret.getValue() = n.getNode() and
         n.refersTo(_, result, _)
     )
 }
 
-from ClassObject t, FunctionObject iter
-where exists(ClassObject ret_t | iter = iter_method(t) and
-          ret_t = return_type(iter) and
-          not ret_t.isIterator()
-      )
-
-select iter, "The '__iter__' method of iterable class $@ does not return an iterator.", t, t.getName()
+from ClassObject iterable, FunctionObject iter, ClassObject iterator
+where
+    iter = iterable.lookupAttribute("__iter__") and
+    iterator = return_type(iter) and
+    not iterator.isIterator()
+select iterator,
+    "Class " + iterator.getName() +
+        " is returned as an iterator (by $@) but does not fully implement the iterator interface.",
+    iter, iter.getName()