Python: Experiments with nested comprhensions

Author: yoff

python/ql/src/experimental/dataflow/internal/DataFlowPrivate.qll

@@ -435,16 +435,26 @@ predicate comprehensionReadStep(CfgNode nodeFrom, Content c, EssaNode nodeTo) {
   //   nodeFrom is `l`, cfg node
   //   nodeTo is `x`, essa var
   //   c denotes element of list or set
-  exists(For f, Comp comp |
-    f = getCompFor(comp) and
-    nodeFrom.getNode().getNode() = getCompIter(comp) and
-    nodeTo.getVar().getDefinition().(AssignmentDefinition).getDefiningNode().getNode() =
-      f.getTarget() and
-    (
-      c instanceof ListElementContent
-      or
-      c instanceof SetElementContent
+  exists(Comp comp |
+    // outermost for
+    exists(For f |
+      f = getCompFor(comp) and
+      nodeFrom.getNode().getNode() = getCompIter(comp) and
+      nodeTo.getVar().getDefinition().(AssignmentDefinition).getDefiningNode().getNode() =
+        f.getTarget()
+    )
+    or
+    // an inner for
+    exists(int n |
+      nodeFrom.getNode().getNode() = comp.getNthInnerLoop(n + 1).getIter() and
+      nodeTo.getVar().getDefinition().(AssignmentDefinition).getDefiningNode().getNode() =
+        comp.getNthInnerLoop(n).getTarget()
     )
+  ) and
+  (
+    c instanceof ListElementContent
+    or
+    c instanceof SetElementContent
   )
 }
 

python/ql/src/semmle/python/Comprehensions.qll

@@ -4,12 +4,15 @@ import python
 abstract class Comp extends Expr {
   abstract Function getFunction();
 
+  /** Gets the iterable of this set comprehension. */
+  abstract Expr getIterable();
+
   /** Gets the iteration variable for the nth innermost generator of this list comprehension */
   Variable getIterationVariable(int n) {
     result.getAnAccess() = this.getNthInnerLoop(n).getTarget()
   }
 
-  private For getNthInnerLoop(int n) {
+  For getNthInnerLoop(int n) {
     n = 0 and result = this.getFunction().getStmt(0)
     or
     result = this.getNthInnerLoop(n - 1).getStmt(0)
@@ -62,6 +65,8 @@ class ListComp extends ListComp_, Comp {
 
   override Function getFunction() { result = ListComp_.super.getFunction() }
 
+  override Expr getIterable() { result = ListComp_.super.getIterable() }
+
   override string toString() { result = ListComp_.super.toString() }
 
   override Expr getElt() { result = Comp.super.getElt() }
@@ -79,6 +84,8 @@ class SetComp extends SetComp_, Comp {
   override predicate hasSideEffects() { any() }
 
   override Function getFunction() { result = SetComp_.super.getFunction() }
+
+  override Expr getIterable() { result = SetComp_.super.getIterable() }
 }
 
 /** A dictionary comprehension, such as `{ k:v for k, v in enumerate("0123456789") }` */
@@ -93,6 +100,8 @@ class DictComp extends DictComp_, Comp {
   override predicate hasSideEffects() { any() }
 
   override Function getFunction() { result = DictComp_.super.getFunction() }
+
+  override Expr getIterable() { result = DictComp_.super.getIterable() }
 }
 
 /** A generator expression, such as `(var for var in iterable)` */
@@ -107,4 +116,6 @@ class GeneratorExp extends GeneratorExp_, Comp {
   override predicate hasSideEffects() { any() }
 
   override Function getFunction() { result = GeneratorExp_.super.getFunction() }
+
+  override Expr getIterable() { result = GeneratorExp_.super.getIterable() }
 }