Python: refactor awaited

python/ql/lib/semmle/python/ApiGraphs.qll

@@ -570,8 +570,7 @@ module API {
      * API graph node for the prefix `foo`), in accordance with the usual semantics of Python.
      */
 
-    private import semmle.python.internal.Awaited
-
+    // private import semmle.python.internal.Awaited
     cached
     newtype TApiNode =
       /** The root of the API graph. */
@@ -762,11 +761,8 @@ module API {
         )
         or
         // awaiting
-        exists(DataFlow::Node awaitedValue |
-          lbl = Label::await() and
-          ref = awaited(awaitedValue) and
-          pred.flowsTo(awaitedValue)
-        )
+        lbl = Label::await() and
+        ref = pred.getAnAwaited()
       )
       or
       exists(DataFlow::Node def, PY::CallableExpr fn |

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

@@ -10,6 +10,7 @@ private import python
 import DataFlowPublic
 private import DataFlowPrivate
 private import semmle.python.internal.CachedStages
+private import semmle.python.internal.Awaited
 
 /**
  * A data flow node that is a source of local flow. This includes things like
@@ -95,6 +96,11 @@ class LocalSourceNode extends Node {
    */
   CallCfgNode getACall() { Cached::call(this, result) }
 
+  /**
+   * Gets an awaited value from this node.
+   */
+  Node getAnAwaited() { Cached::await(this, result) }
+
   /**
    * Gets a call to the method `methodName` on this node.
    *
@@ -225,4 +231,15 @@ private module Cached {
       n = call.getFunction()
     )
   }
+
+  /**
+   * Holds if `node` flows to the awaited value of `awaited`.
+   */
+  cached
+  predicate await(LocalSourceNode node, Node awaited) {
+    exists(Node awaitedValue |
+      node.flowsTo(awaitedValue) and
+      awaited = awaited(awaitedValue)
+    )
+  }
 }