Python: Add missing flow for AssignmentExpr nodes

Also extend the tests surrounding this construct to be a bit more comprehensive.

Co-authored-by: Rasmus Lerchedahl Petersen <yoff@github.com>

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

@@ -547,6 +547,31 @@ class IfExprNode extends ControlFlowNode {
   override IfExp getNode() { result = super.getNode() }
 }
 
+/** A control flow node corresponding to an assignment expression such as `lhs := rhs`. */
+class AssignmentExprNode extends ControlFlowNode {
+  AssignmentExprNode() { toAst(this) instanceof AssignExpr }
+
+  /** Gets the flow node corresponding to the left-hand side of the assignment expression */
+  ControlFlowNode getTarget() {
+    exists(AssignExpr a |
+      this.getNode() = a and
+      a.getTarget() = result.getNode() and
+      result.getBasicBlock().dominates(this.getBasicBlock())
+    )
+  }
+
+  /** Gets the flow node corresponding to the right-hand side of the assignment expression */
+  ControlFlowNode getValue() {
+    exists(AssignExpr a |
+      this.getNode() = a and
+      a.getValue() = result.getNode() and
+      result.getBasicBlock().dominates(this.getBasicBlock())
+    )
+  }
+
+  override AssignExpr getNode() { result = super.getNode() }
+}
+
 /** A control flow node corresponding to a binary expression, such as `x + y` */
 class BinaryExprNode extends ControlFlowNode {
   BinaryExprNode() { toAst(this) instanceof BinaryExpr }
@@ -630,6 +655,8 @@ class DefinitionNode extends ControlFlowNode {
     Stages::AST::ref() and
     exists(Assign a | a.getATarget().getAFlowNode() = this)
     or
+    exists(AssignExpr a | a.getTarget().getAFlowNode() = this)
+    or
     exists(AnnAssign a | a.getTarget().getAFlowNode() = this and exists(a.getValue()))
     or
     exists(Alias a | a.getAsname().getAFlowNode() = this)
@@ -787,6 +814,9 @@ private AstNode assigned_value(Expr lhs) {
   /* lhs = result */
   exists(Assign a | a.getATarget() = lhs and result = a.getValue())
   or
+  /* lhs := result */
+  exists(AssignExpr a | a.getTarget() = lhs and result = a.getValue())
+  or
   /* lhs : annotation = result */
   exists(AnnAssign a | a.getTarget() = lhs and result = a.getValue())
   or

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

@@ -357,6 +357,9 @@ module EssaFlow {
     // If expressions
     nodeFrom.asCfgNode() = nodeTo.asCfgNode().(IfExprNode).getAnOperand()
     or
+    // Assignment expressions
+    nodeFrom.asCfgNode() = nodeTo.asCfgNode().(AssignmentExprNode).getValue()
+    or
     // boolean inline expressions such as `x or y` or `x and y`
     nodeFrom.asCfgNode() = nodeTo.asCfgNode().(BoolExprNode).getAnOperand()
     or

python/ql/test/experimental/dataflow/coverage/test.py

@@ -435,10 +435,14 @@ def test_and(x = True):
 
 
 # 6.12. Assignment expressions
-def test_assignment_expression():
+def test_assignment_expression_flow_lhs():
     x = NONSOURCE
-    SINK(x := SOURCE) #$ MISSING:flow="SOURCE -> x"
+    if x := SOURCE:
+        SINK(x) #$ flow="SOURCE, l:-1 -> x"
 
+def test_assignment_expression_flow_out():
+    x = NONSOURCE
+    SINK(x := SOURCE) #$ flow="SOURCE -> AssignExpr"
 
 # 6.13. Conditional expressions
 def test_conditional_true():
@@ -460,13 +464,13 @@ def test_conditional_false_guards():
 # Condition is evaluated first, so x is SOURCE once chosen
 def test_conditional_evaluation_true():
     x = NONSOURCE
-    SINK(x if (SOURCE == (x := SOURCE)) else NONSOURCE) #$ MISSING:flow="SOURCE -> IfExp"
+    SINK(x if (SOURCE == (x := SOURCE)) else NONSOURCE) #$ flow="SOURCE -> IfExp"
 
 
 # Condition is evaluated first, so x is SOURCE once chosen
 def test_conditional_evaluation_false():
     x = NONSOURCE
-    SINK(NONSOURCE if (NONSOURCE == (x := SOURCE)) else x) #$ MISSING:flow="SOURCE -> IfExp"
+    SINK(NONSOURCE if (NONSOURCE == (x := SOURCE)) else x) #$ flow="SOURCE -> IfExp"
 
 
 # 6.14. Lambdas