Fix handling for some wrappers + add test case

Author: joefarebrother

python/ql/src/Functions/IterReturnsNonSelf.ql

@@ -50,7 +50,7 @@ predicate iterWrapperMethods(Function iter, Function next) {
   exists(string field |
     exists(Return r, DataFlow::Node self, DataFlow::AttrRead read |
       r.getScope() = iter and
-      r.getValue() = iterCall(read).asExpr() and
+      r.getValue() = [iterCall(read).asExpr(), read.asExpr()] and
       read.accesses(self, field) and
       isSelfVar(iter, self.asExpr())
     ) and
@@ -63,15 +63,13 @@ predicate iterWrapperMethods(Function iter, Function next) {
   )
 }
 
-/** Gets a call to `iter(arg)`, `arg.__iter__()`, or `arg` itself (which we assume may already be an iterator). */
+/** Gets a call to `iter(arg)` or `arg.__iter__()`. */
 private DataFlow::CallCfgNode iterCall(DataFlow::Node arg) {
   result.(DataFlow::MethodCallNode).calls(arg, "__iter__")
   or
   result = API::builtin("iter").getACall() and
   arg = result.getArg(0) and
   not exists(result.getArg(1))
-  or
-  result = arg // assume the wrapping field is already an iterator
 }
 
 /** Gets a call to `next(arg)` or `arg.__next__()`. */

python/ql/test/query-tests/Functions/iterators/IterReturnsNonSelf.expected

@@ -1,2 +1,2 @@
 | test.py:5:5:5:23 | Function __iter__ | Iter method of iterator $@ does not return `self`. | test.py:1:1:1:11 | Class Bad1 | Bad1 |
-| test.py:41:5:41:23 | Function __iter__ | Iter method of iterator $@ does not return `self`. | test.py:32:1:32:21 | Class FalsePositive1 | FalsePositive1 |
+| test.py:51:5:51:23 | Function __iter__ | Iter method of iterator $@ does not return `self`. | test.py:42:1:42:21 | Class FalsePositive1 | FalsePositive1 |

python/ql/test/query-tests/Functions/iterators/test.py

@@ -21,8 +21,18 @@ def __next__(self):
 
     def __iter__(self): # GOOD: iter and next are wrappers around a field
         return self._it.__iter__()
-
+    
 class Good3:
+    def __init__(self):
+        self._it = iter([0,0,0])
+
+    def __next__(self):
+        return self._it.__next__()
+
+    def __iter__(self): # GOOD: iter and next are wrappers around a field
+        return self._it
+
+class Good4:
     def __next__(self):
         return 0