Python: Add API graph support for parameter annotations

Adds API graph support for observing that in
```python
def foo(x : Bar): ...
```
The variable `x` is likely to be an instance of the type `Bar` inside
this function.
In particular, we add `getInstanceFromAnnotation` as a predicate on API
graph nodes that tracks this step (corresponding to a new edge type
labeled with "annotation" in the API graph), and extend the existing
`getAnInstance` predicate to also include instances arising from type
annotations.

A more complete solution would also add support for annotated
assignments (`x : Foo = ...` or just `x : Foo`) as well as track types
through type aliases (`type Foo = Bar`). This turns out to be
non-trivial, however, as these type constructs don't have any CFG nodes
(and so no data-flow nodes by default either). In order to not have
perfect be the enemy of good, this commit is only targeting the type
parameter case (which is also likely to be the most common use case
anyway).

The tests for API graphs have been extended accordingly, including tests
for the kinds of type ascriptions that we _don't_ currently model in API
graphs (marked with `MISSING:` in the inline tests).

Author: tausbn

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()