Merge pull request #5944 from RasmusWL/async-api-graph-tests

Approved by tausbn

python/ql/src/semmle/python/ApiGraphs.qll

@@ -97,6 +97,11 @@ module API {
      */
     Node getASubclass() { result = getASuccessor(Label::subclass()) }
 
+    /**
+     * Gets a node representing the result from awaiting this node.
+     */
+    Node getAwaited() { result = getASuccessor(Label::await()) }
+
     /**
      * Gets a string representation of the lexicographically least among all shortest access paths
      * from the root to this node.
@@ -469,6 +474,14 @@ module API {
         exists(DataFlow::Node superclass | pred.flowsTo(superclass) |
           ref.asExpr().(ClassExpr).getABase() = superclass.asExpr()
         )
+        or
+        // awaiting
+        exists(Await await, DataFlow::Node awaitedValue |
+          lbl = Label::await() and
+          ref.asExpr() = await and
+          await.getValue() = awaitedValue.asExpr() and
+          pred.flowsTo(awaitedValue)
+        )
       )
       or
       // Built-ins, treated as members of the module `builtins`
@@ -585,5 +598,9 @@ private module Label {
   /** Gets the `return` edge label. */
   string return() { result = "getReturn()" }
 
+  /** Gets the `subclass` edge label. */
   string subclass() { result = "getASubclass()" }
+
+  /** Gets the `await` edge label. */
+  string await() { result = "getAwaited()" }
 }

python/ql/test/experimental/dataflow/ApiGraphs/async_test.py

@@ -0,0 +1,17 @@
+import pkg # $ use=moduleImport("pkg")
+
+async def foo():
+    coro = pkg.async_func() # $ use=moduleImport("pkg").getMember("async_func").getReturn()
+    coro # $ use=moduleImport("pkg").getMember("async_func").getReturn()
+    result = await coro # $ use=moduleImport("pkg").getMember("async_func").getReturn().getAwaited()
+    result # $ use=moduleImport("pkg").getMember("async_func").getReturn().getAwaited()
+    return result # $ use=moduleImport("pkg").getMember("async_func").getReturn().getAwaited()
+
+async def bar():
+    result = await pkg.async_func() # $ use=moduleImport("pkg").getMember("async_func").getReturn().getAwaited()
+    return result # $ use=moduleImport("pkg").getMember("async_func").getReturn().getAwaited()
+
+def check_annotations():
+    # Just to make sure how annotations should look like :)
+    result = pkg.sync_func() # $ use=moduleImport("pkg").getMember("sync_func").getReturn()
+    return result # $ use=moduleImport("pkg").getMember("sync_func").getReturn()

python/ql/test/experimental/dataflow/ApiGraphs/options

@@ -1 +1 @@
-semmle-extractor-options: --lang=3
+semmle-extractor-options: --lang=3 --max-import-depth=1

python/ql/test/experimental/dataflow/ApiGraphs/use.ql

@@ -13,7 +13,8 @@ class ApiUseTest extends InlineExpectationsTest {
     l = n.getLocation() and
     // Module variable nodes have no suitable location, so it's best to simply exclude them entirely
     // from the inline tests.
-    not n instanceof DataFlow::ModuleVariableNode
+    not n instanceof DataFlow::ModuleVariableNode and
+    exists(l.getFile().getRelativePath())
   }
 
   override predicate hasActualResult(Location location, string element, string tag, string value) {