Merge pull request #14851 from RasmusWL/variable-caputre-list-comprehension

Python: Add test for variable reference in list comprehension

Author: RasmusWL

python/ql/test/experimental/dataflow/coverage/localFlow.expected

@@ -5,13 +5,13 @@
 | test.py:42:10:42:26 | ControlFlowNode for Tuple | test.py:42:5:42:5 | SSA variable x |
 | test.py:43:5:43:5 | SSA variable y | test.py:44:10:44:10 | ControlFlowNode for y |
 | test.py:43:9:43:12 | ControlFlowNode for Subscript | test.py:43:5:43:5 | SSA variable y |
-| test.py:187:1:187:53 | GSSA Variable SINK | test.py:189:5:189:8 | ControlFlowNode for SINK |
-| test.py:187:1:187:53 | GSSA Variable SOURCE | test.py:188:25:188:30 | ControlFlowNode for SOURCE |
-| test.py:188:5:188:5 | SSA variable x | test.py:189:10:189:10 | ControlFlowNode for x |
-| test.py:188:9:188:68 | ControlFlowNode for .0 | test.py:188:9:188:68 | SSA variable .0 |
-| test.py:188:9:188:68 | ControlFlowNode for ListComp | test.py:188:5:188:5 | SSA variable x |
-| test.py:188:9:188:68 | SSA variable .0 | test.py:188:9:188:68 | ControlFlowNode for .0 |
-| test.py:188:16:188:16 | SSA variable v | test.py:188:45:188:45 | ControlFlowNode for v |
-| test.py:188:40:188:40 | SSA variable u | test.py:188:56:188:56 | ControlFlowNode for u |
-| test.py:188:51:188:51 | SSA variable z | test.py:188:67:188:67 | ControlFlowNode for z |
-| test.py:188:62:188:62 | SSA variable y | test.py:188:10:188:10 | ControlFlowNode for y |
+| test.py:208:1:208:53 | GSSA Variable SINK | test.py:210:5:210:8 | ControlFlowNode for SINK |
+| test.py:208:1:208:53 | GSSA Variable SOURCE | test.py:209:25:209:30 | ControlFlowNode for SOURCE |
+| test.py:209:5:209:5 | SSA variable x | test.py:210:10:210:10 | ControlFlowNode for x |
+| test.py:209:9:209:68 | ControlFlowNode for .0 | test.py:209:9:209:68 | SSA variable .0 |
+| test.py:209:9:209:68 | ControlFlowNode for ListComp | test.py:209:5:209:5 | SSA variable x |
+| test.py:209:9:209:68 | SSA variable .0 | test.py:209:9:209:68 | ControlFlowNode for .0 |
+| test.py:209:16:209:16 | SSA variable v | test.py:209:45:209:45 | ControlFlowNode for v |
+| test.py:209:40:209:40 | SSA variable u | test.py:209:56:209:56 | ControlFlowNode for u |
+| test.py:209:51:209:51 | SSA variable z | test.py:209:67:209:67 | ControlFlowNode for z |
+| test.py:209:62:209:62 | SSA variable y | test.py:209:10:209:10 | ControlFlowNode for y |

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

@@ -120,6 +120,27 @@ def test_nested_list_display():
     SINK(x[0]) #$ MISSING:flow="SOURCE, l:-1 -> x[0]"
 
 
+@expects(6)
+def test_list_comprehension_with_tuple_result():
+    # confirms that this simple case works as expected
+    t = (SOURCE, NONSOURCE)
+    SINK(t[0]) # $ flow="SOURCE, l:-1 -> t[0]"
+    SINK_F(t[1])
+
+    # confirms that this simple case works as expected
+    l1 = [(SOURCE, NONSOURCE) for _ in [1]]
+    SINK(l1[0][0]) # $ flow="SOURCE, l:-1 -> l1[0][0]"
+    SINK_F(l1[0][1])
+
+    # stops working when using reference to variable, due to variables being captured by
+    # the function we internally generate for the comprehension body.
+    s = SOURCE
+    ns = NONSOURCE
+    l3 = [(s, ns) for _ in [1]]
+    SINK(l3[0][0]) # $ MISSING: flow="SOURCE, l:-3 -> l3[0][0]"
+    SINK_F(l3[0][1])
+
+
 # 6.2.6. Set displays
 def test_set_display():
     x = {SOURCE}