Merge pull request #18112 from github/tausbn/add-api-graph-support-for-parameter-annotations

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

@@ -195,6 +195,12 @@ module API {
      */
     Node getReturn() { result = this.getASuccessor(Label::return()) }
 
+    /**
+     * Gets a node representing instances of the class represented by this node, as specified via
+     * type annotations.
+     */
+    Node getInstanceFromAnnotation() { result = this.getASuccessor(Label::annotation()) }
+
     /**
      * Gets a node representing the `i`th parameter of the function represented by this node.
      *
@@ -229,7 +235,9 @@ module API {
     /**
      * Gets a node representing an instance of the class (or a transitive subclass of the class) represented by this node.
      */
-    Node getAnInstance() { result = this.getASubclass*().getReturn() }
+    Node getAnInstance() {
+      result in [this.getASubclass*().getReturn(), this.getASubclass*().getInstanceFromAnnotation()]
+    }
 
     /**
      * Gets a node representing the result from awaiting this node.
@@ -834,6 +842,10 @@ module API {
         lbl = Label::return() and
         ref = pred.getACall()
         or
+        // Getting an instance via a type annotation
+        lbl = Label::annotation() and
+        ref = pred.getAnAnnotatedInstance()
+        or
         // Awaiting a node that is a use of `base`
         lbl = Label::await() and
         ref = pred.getAnAwaited()
@@ -1079,6 +1091,7 @@ module API {
         } or
         MkLabelSelfParameter() or
         MkLabelReturn() or
+        MkLabelAnnotation() or
         MkLabelSubclass() or
         MkLabelAwait() or
         MkLabelSubscript() or
@@ -1148,6 +1161,11 @@ module API {
         override string toString() { result = "getReturn()" }
       }
 
+      /** A label for annotations. */
+      class LabelAnnotation extends ApiLabel, MkLabelAnnotation {
+        override string toString() { result = "getAnnotatedInstance()" }
+      }
+
       /** A label that gets the subclass of a class. */
       class LabelSubclass extends ApiLabel, MkLabelSubclass {
         override string toString() { result = "getASubclass()" }
@@ -1207,6 +1225,9 @@ module API {
     /** Gets the `return` edge label. */
     LabelReturn return() { any() }
 
+    /** Gets the `annotation` edge label. */
+    LabelAnnotation annotation() { any() }
+
     /** Gets the `subclass` edge label. */
     LabelSubclass subclass() { any() }
 

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

@@ -119,6 +119,11 @@ class LocalSourceNode extends Node {
    */
   CallCfgNode getACall() { Cached::call(this, result) }
 
+  /**
+   * Gets a node that has this node as its annotation.
+   */
+  Node getAnAnnotatedInstance() { Cached::annotatedInstance(this, result) }
+
   /**
    * Gets an awaited value from this node.
    */
@@ -275,6 +280,17 @@ private module Cached {
     )
   }
 
+  cached
+  predicate annotatedInstance(LocalSourceNode node, Node instance) {
+    exists(ExprNode n | node.flowsTo(n) |
+      instance.asCfgNode().getNode() =
+        any(AnnAssign ann | ann.getAnnotation() = n.asExpr()).getTarget()
+      or
+      instance.asCfgNode().getNode() =
+        any(Parameter p | p.getAnnotation() = n.asCfgNode().getNode())
+    )
+  }
+
   /**
    * Holds if `node` flows to a value that, when awaited, results in `awaited`.
    */