Python: Add API graph support for parameter annotations

python/ql/lib/change-notes/2024-11-26-parameter-annotation-api-graph-support.md

@@ -0,0 +1,5 @@
+---
+category: feature
+---
+
+- Added support for parameter annotations in API graphs. This means that in a function definition such as `def foo(x: Bar): ...`, you can now use the `getInstanceFromAnnotation()` method to step from `Bar` to `x`. In addition to this, the `getAnInstance` method now also includes instances arising from 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`.
    */

python/ql/test/library-tests/ApiGraphs/py3/test_annotations.py

@@ -0,0 +1,25 @@
+from types import AssignmentAnnotation, ParameterAnnotation
+
+def test_annotated_assignment():
+    local_x : AssignmentAnnotation = create_x() #$ MISSING: use=moduleImport("types").getMember("AssignmentAnnotation")
+    local_x #$ MISSING: use=moduleImport("types").getMember("AssignmentAnnotation").getAnnotatedInstance()
+
+global_x : AssignmentAnnotation #$ use=moduleImport("types").getMember("AssignmentAnnotation")
+global_x #$ MISSING: use=moduleImport("types").getMember("AssignmentAnnotation").getAnnotatedInstance()
+
+def test_parameter_annotation(parameter_y: ParameterAnnotation): #$ use=moduleImport("types").getMember("ParameterAnnotation")
+    parameter_y #$ use=moduleImport("types").getMember("ParameterAnnotation").getAnnotatedInstance()
+
+type Alias = AssignmentAnnotation
+
+global_z : Alias #$ MISSING: use=moduleImport("types").getMember("AssignmentAnnotation")
+global_z #$ MISSING: use=moduleImport("types").getMember("AssignmentAnnotation").getAnnotatedInstance()
+
+def test_parameter_alias(parameter_z: Alias): #$ MISSING: use=moduleImport("types").getMember("AssignmentAnnotation")
+    parameter_z #$ MISSING: use=moduleImport("types").getMember("AssignmentAnnotation").getAnnotatedInstance()
+
+# local type aliases
+def test_local_type_alias():
+    type LocalAlias = AssignmentAnnotation
+    local_alias : LocalAlias = create_value() #$ MISSING: use=moduleImport("types").getMember("AssignmentAnnotation")
+    local_alias #$ MISSING: use=moduleImport("types").getMember("AssignmentAnnotation").getAnnotatedInstance()