Python: DefinitionNode for with

Author: yoff

python/ql/lib/semmle/python/Flow.qll

@@ -672,6 +672,8 @@ class DefinitionNode extends ControlFlowNode {
     exists(For for | for.getTarget().getAFlowNode() = this)
     or
     exists(Parameter param | this = param.asName().getAFlowNode() and exists(param.getDefault()))
+    or
+    exists(With with | this = with.getOptionalVars().getAFlowNode())
   }
 
   /** flow node corresponding to the value assigned for the definition corresponding to this flow node */
@@ -840,7 +842,11 @@ private AstNode assigned_value(Expr lhs) {
   /* for lhs in seq: => `result` is the `for` node, representing the `iter(next(seq))` operation. */
   result.(For).getTarget() = lhs
   or
+  // def f(a = 42):  =>  `result` is 42
   exists(Parameter param | lhs = param.asName() and result = param.getDefault())
+  or
+  // with f(42) as x:  =>  `result` is `f(42)`
+  exists(With with | lhs = with.getOptionalVars() and result = with.getContextExpr())
 }
 
 predicate nested_sequence_assign(