Python: Add syntactic support for yield in contextlib.contextmanager

Author: RasmusWL

python/ql/lib/change-notes/2023-10-17-contextmanager.md

@@ -0,0 +1,4 @@
+---
+category: minorAnalysis
+---
+* Added support for functions decorated with `contextlib.contextmanager`.

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

@@ -240,6 +240,19 @@ predicate hasPropertyDecorator(Function func) {
   )
 }
 
+/**
+ * Holds if the function `func` has a `contextlib.contextmanager`.
+ */
+predicate hasContextmanagerDecorator(Function func) {
+  exists(ControlFlowNode contextmanager |
+    contextmanager.(NameNode).getId() = "contextmanager" and contextmanager.(NameNode).isGlobal()
+    or
+    contextmanager.(AttrNode).getObject("contextmanager").(NameNode).getId() = "contextlib"
+  |
+    func.getADecorator() = contextmanager.getNode()
+  )
+}
+
 // =============================================================================
 // Callables
 // =============================================================================
@@ -1604,6 +1617,16 @@ class ExtractedReturnNode extends ReturnNode, CfgNode {
   override ReturnKind getKind() { any() }
 }
 
+/** A data flow node that represents a value returned by a callable. */
+class YieldNodeInContextManagerFunction extends ReturnNode, CfgNode {
+  YieldNodeInContextManagerFunction() {
+    hasContextmanagerDecorator(node.getScope()) and
+    node = any(Yield yield).getValue().getAFlowNode()
+  }
+
+  override ReturnKind getKind() { any() }
+}
+
 /** A data-flow node that represents the output of a call. */
 abstract class OutNode extends Node {
   /** Gets the underlying call, where this node is a corresponding output of kind `kind`. */

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

@@ -223,14 +223,14 @@ def managed_resource():
     yield x # $ tracked
 
 def test_context_manager():
-    with managed_resource() as x: # $ MISSING: tracked
-        print(x) # $ MISSING: tracked
+    with managed_resource() as x: # $ tracked
+        print(x) # $ tracked
 
 @contextlib.contextmanager
 def managed_resource2():
     x = tracked # $ tracked
     yield x # $ tracked
 
 def test_context_manager2():
-    with managed_resource2() as x: # $ MISSING: tracked
-        print(x) # $ MISSING: tracked
+    with managed_resource2() as x: # $ tracked
+        print(x) # $ tracked