Python: flow through yield

- add yield as a dataflow return
- replace comprehension store step
   with a store step to the yield

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

@@ -1755,6 +1755,8 @@ class SynthCompCapturedVariablesArgumentNode extends Node, TSynthCompCapturedVar
   override Scope getScope() { result = comp.getFunction() }
 
   override Location getLocation() { result = comp.getLocation() }
+
+  Comp getComprehension() { result = comp }
 }
 
 /** Gets a viable run-time target for the call `call`. */
@@ -1796,7 +1798,10 @@ abstract class ReturnNode extends Node {
 /** A data flow node that represents a value returned by a callable. */
 class ExtractedReturnNode extends ReturnNode, CfgNode {
   // See `TaintTrackingImplementation::returnFlowStep`
-  ExtractedReturnNode() { node = any(Return ret).getValue().getAFlowNode() }
+  ExtractedReturnNode() {
+    node = any(Return ret).getValue().getAFlowNode() or
+    node = any(Yield yield).getAFlowNode()
+  }
 
   override ReturnKind getKind() { any() }
 }

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

@@ -168,6 +168,37 @@ private predicate synthDictSplatArgumentNodeStoreStep(
   )
 }
 
+private predicate yieldStoreStep(Node nodeFrom, Content c, Node nodeTo) {
+  exists(Yield yield, Function func |
+    nodeTo.asCfgNode() = yield.getAFlowNode() and
+    nodeFrom.asCfgNode() = yield.getValue().getAFlowNode() and
+    func.containsInScope(yield)
+  |
+    exists(Comp comp | func = comp.getFunction() |
+      (
+        comp instanceof ListComp or
+        comp instanceof GeneratorExp
+      ) and
+      c instanceof ListElementContent
+      or
+      comp instanceof SetComp and
+      c instanceof SetElementContent
+      or
+      comp instanceof DictComp and
+      c instanceof DictionaryElementContent
+    )
+    or
+    not exists(Comp comp | func = comp.getFunction()) and
+    (
+      c instanceof ListElementContent
+      or
+      c instanceof SetElementContent
+      or
+      c instanceof DictionaryElementContent
+    )
+  )
+}
+
 /**
  * Ensures that the a `**kwargs` parameter will not contain elements with names of
  * keyword parameters.
@@ -668,8 +699,6 @@ predicate storeStep(Node nodeFrom, ContentSet c, Node nodeTo) {
   or
   setStoreStep(nodeFrom, c, nodeTo)
   or
-  comprehensionStoreStep(nodeFrom, c, nodeTo)
-  or
   attributeStoreStep(nodeFrom, c, nodeTo)
   or
   matchStoreStep(nodeFrom, c, nodeTo)
@@ -683,6 +712,8 @@ predicate storeStep(Node nodeFrom, ContentSet c, Node nodeTo) {
   or
   synthDictSplatArgumentNodeStoreStep(nodeFrom, c, nodeTo)
   or
+  yieldStoreStep(nodeFrom, c, nodeTo)
+  or
   VariableCapture::storeStep(nodeFrom, c, nodeTo)
 }