Merge pull request #3040 from BekaValentine/python-objectapi-to-valueapi-iterreturnsnonself

Python: ObjectAPI to ValueAPI: IterReturnsNonSelf

Author: RasmusWL

python/ql/src/Functions/IterReturnsNonSelf.ql

@@ -12,9 +12,7 @@
 
 import python
 
-Function iter_method(ClassObject t) {
-    result = t.lookupAttribute("__iter__").(FunctionObject).getFunction()
-}
+Function iter_method(ClassValue t) { result = t.lookup("__iter__").(FunctionValue).getScope() }
 
 predicate is_self(Name value, Function f) { value.getVariable() = f.getArg(0).(Name).getVariable() }
 
@@ -26,7 +24,7 @@ predicate returns_non_self(Function f) {
     exists(Return r | r.getScope() = f and not exists(r.getValue()))
 }
 
-from ClassObject t, Function iter
+from ClassValue t, Function iter
 where t.isIterator() and iter = iter_method(t) and returns_non_self(iter)
 select t, "Class " + t.getName() + " is an iterator but its $@ method does not return 'self'.",
     iter, iter.getName()

python/ql/src/semmle/python/objects/ObjectAPI.qll

@@ -430,6 +430,29 @@ class ClassValue extends Value {
         this.hasAttribute("__getitem__")
     }
 
+    /** Holds if this class is an iterator. */
+    predicate isIterator() {
+        this.hasAttribute("__iter__") and
+        (
+            major_version() = 3 and this.hasAttribute("__next__")
+            or
+            /*
+             * Because 'next' is a common method name we need to check that an __iter__
+             * method actually returns this class. This is not needed for Py3 as the
+             * '__next__' method exists to define a class as an iterator.
+             */
+
+            major_version() = 2 and
+            this.hasAttribute("next") and
+            exists(ClassValue other, FunctionValue iter | other.declaredAttribute("__iter__") = iter |
+                iter.getAnInferredReturnType() = this
+            )
+        )
+        or
+        /* This will be redundant when we have C class information */
+        this = ClassValue::generator()
+    }
+
     /** Holds if this class is a container(). That is, does it have a __getitem__ method. */
     predicate isContainer() { exists(this.lookup("__getitem__")) }
 
@@ -583,11 +606,7 @@ abstract class FunctionValue extends CallableValue {
     }
 
     /** Gets a class that this function may return */
-    ClassValue getAnInferredReturnType() {
-        result = TBuiltinClassObject(this.(BuiltinFunctionObjectInternal).getReturnType())
-        or
-        result = TBuiltinClassObject(this.(BuiltinMethodObjectInternal).getReturnType())
-    }
+    abstract ClassValue getAnInferredReturnType();
 }
 
 /** Class representing Python functions */
@@ -616,6 +635,13 @@ class PythonFunctionValue extends FunctionValue {
 
     /** Gets a control flow node corresponding to a return statement in this function */
     ControlFlowNode getAReturnedNode() { result = this.getScope().getAReturnValueFlowNode() }
+
+    override ClassValue getAnInferredReturnType() {
+        /* We have to do a special version of this because builtin functions have no
+         * explicit return nodes that we can query and get the class of.
+         */
+        result = this.getAReturnedNode().pointsTo().getClass()
+    }
 }
 
 /** Class representing builtin functions, such as `len` or `print` */
@@ -627,6 +653,13 @@ class BuiltinFunctionValue extends FunctionValue {
     override int minParameters() { none() }
 
     override int maxParameters() { none() }
+
+    override ClassValue getAnInferredReturnType() {
+        /* We have to do a special version of this because builtin functions have no
+         * explicit return nodes that we can query and get the class of.
+         */
+        result = TBuiltinClassObject(this.(BuiltinFunctionObjectInternal).getReturnType())
+    }
 }
 
 /** Class representing builtin methods, such as `list.append` or `set.add` */
@@ -644,6 +677,10 @@ class BuiltinMethodValue extends FunctionValue {
     override int minParameters() { none() }
 
     override int maxParameters() { none() }
+
+    override ClassValue getAnInferredReturnType() {
+        result = TBuiltinClassObject(this.(BuiltinMethodObjectInternal).getReturnType())
+    }
 }
 
 /**