Python: Add getAwaited to API::Node

I _really_ wanted to call this `.await()`, but that did not fit in with
the convention, or the corresponding `getPromised` in JS.

https://github.com/github/codeql/blob/54f191cfe37e136bcc7189b2fb01f44dbb01758b/javascript/ql/src/semmle/javascript/ApiGraphs.qll#L184

Author: RasmusWL

python/change-notes/2021-05-21-api-graph-await.md

@@ -0,0 +1,2 @@
+lgtm,codescanning
+* API graph nodes now contain a `getAwaited()` member predicate, for getting the result of awaiting an item, such as `await foo`.

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`
@@ -587,4 +600,7 @@ private module Label {
 
   /** 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

@@ -3,13 +3,13 @@
 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()
-    result # $ MISSING: use=...
-    return result # $ MISSING: use=...
+    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()
-    return result # $ MISSING: use=...
+    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 :)