python: capture flow through comprehensions

- add comprehension functions as `DataFlowCallable`s
- add comprehension call as `DataFlowCall`
- create capture argument node for comprehension calls

Author: yoff

python/ql/lib/semmle/python/dataflow/new/internal/DataFlowDispatch.qll

@@ -320,6 +320,9 @@ newtype TDataFlowCallable =
     // same to keep things easy to reason about (and therefore exclude things that do
     // not have a definition)
     exists(func.getDefinition())
+    or
+    // ...scratch that, variable capture requires a callable
+    exists(Comp c | c.getFunction() = func)
   } or
   /** see QLDoc for `DataFlowModuleScope` for why we need this. */
   TModule(Module m) or
@@ -1382,6 +1385,7 @@ private predicate sameEnclosingCallable(Node node1, Node node2) {
 // =============================================================================
 newtype TDataFlowCall =
   TNormalCall(CallNode call, Function target, CallType type) { resolveCall(call, target, type) } or
+  TComprehensionCall(Comp c) or
   TPotentialLibraryCall(CallNode call) or
   /** A synthesized call inside a summarized callable */
   TSummaryCall(
@@ -1468,6 +1472,30 @@ class NormalCall extends ExtractedDataFlowCall, TNormalCall {
   CallType getCallType() { result = type }
 }
 
+class ComprehensionCall extends ExtractedDataFlowCall, TComprehensionCall {
+  Comp c;
+  Function target;
+
+  ComprehensionCall() {
+    this = TComprehensionCall(c) and
+    target = c.getFunction()
+  }
+
+  Comp getComprehension() { result = c }
+
+  override string toString() { result = "comprehension call" }
+
+  override ControlFlowNode getNode() { result.getNode() = c }
+
+  override Scope getScope() { result = c.getScope() }
+
+  override DataFlowCallable getCallable() { result.(DataFlowFunction).getScope() = target }
+
+  override ArgumentNode getArgument(ArgumentPosition apos) { none() }
+
+  override Location getLocation() { result = c.getLocation() }
+}
+
 /**
  * A potential call to a summarized callable, a `LibraryCallable`.
  *
@@ -1698,6 +1726,24 @@ class CapturedVariablesArgumentNode extends CfgNode, ArgumentNode {
   }
 }
 
+class ComprehensionCapturedVariablesArgumentNode extends CfgNode, ArgumentNode {
+  Comp comp;
+
+  ComprehensionCapturedVariablesArgumentNode() {
+    node.getNode() = comp and
+    exists(Function target | target = comp.getFunction() |
+      target = any(VariableCapture::CapturedVariable v).getACapturingScope()
+    )
+  }
+
+  override string toString() { result = "Capturing closure argument (comp)" }
+
+  override predicate argumentOf(DataFlowCall call, ArgumentPosition pos) {
+    call.(ComprehensionCall).getComprehension() = comp and
+    pos.isLambdaSelf()
+  }
+}
+
 /** Gets a viable run-time target for the call `call`. */
 DataFlowCallable viableCallable(DataFlowCall call) {
   call instanceof ExtractedDataFlowCall and

python/ql/test/library-tests/dataflow/coverage/test.py

@@ -137,7 +137,7 @@ def test_list_comprehension_with_tuple_result():
     s = SOURCE
     ns = NONSOURCE
     l3 = [(s, ns) for _ in [1]]
-    SINK(l3[0][0]) # $ MISSING: flow="SOURCE, l:-3 -> l3[0][0]"
+    SINK(l3[0][0]) # $ flow="SOURCE, l:-3 -> l3[0][0]"
     SINK_F(l3[0][1])