Merge pull request #19672 from github/tausbn/python-support-type-annotations-in-call-graph

Python: Support type annotations in call graph

Author: tausbn

python/ql/lib/change-notes/2025-06-04-call-graph-type-annotations.md

@@ -0,0 +1,4 @@
+---
+category: minorAnalysis
+---
+* Type annotations such as `foo : Bar` are now treated by the call graph as an indication that `foo` may be an instance of `Bar`.

python/ql/lib/semmle/python/Exprs.qll

@@ -762,6 +762,17 @@ class Annotation extends Expr {
     or
     this = any(FunctionExpr f).getReturns()
   }
+
+  /** Gets the expression that this annotation annotates. */
+  Expr getAnnotatedExpression() {
+    result = any(AnnAssign a | a.getAnnotation() = this).getTarget()
+    or
+    result = any(Parameter p | p.getAnnotation() = this)
+    or
+    exists(FunctionExpr f, Return r |
+      this = f.getReturns() and r.getScope() = f.getInnerScope() and result = r.getValue()
+    )
+  }
 }
 
 /* Expression Contexts */

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

@@ -580,6 +580,11 @@ private module TrackClassInstanceInput implements CallGraphConstruction::Simple:
   class State = Class;
 
   predicate start(Node start, Class cls) {
+    exists(Annotation ann |
+      ann = classTracker(cls).asExpr() and
+      start.asExpr() = ann.getAnnotatedExpression()
+    )
+    or
     resolveClassCall(start.(CallCfgNode).asCfgNode(), cls)
     or
     // result of `super().__new__` as used in a `__new__` method implementation

python/ql/test/experimental/library-tests/CallGraph-type-annotations/InlineCallGraphTest.expected

@@ -0,0 +1,6 @@
+testFailures
+debug_callableNotUnique
+pointsTo_found_typeTracker_notFound
+typeTracker_found_pointsTo_notFound
+| type_annotations.py:6:5:6:14 | ControlFlowNode for Attribute() | Foo.method |
+| type_annotations.py:16:5:16:14 | ControlFlowNode for Attribute() | Foo.method |

python/ql/test/experimental/library-tests/CallGraph-type-annotations/InlineCallGraphTest.qlref

@@ -0,0 +1 @@
+../CallGraph/InlineCallGraphTest.ql

python/ql/test/experimental/library-tests/CallGraph-type-annotations/type_annotations.py

@@ -0,0 +1,33 @@
+class Foo:
+    def method(self):
+        pass
+
+def test_parameter_annotation(x: Foo):
+    x.method() #$ tt=Foo.method
+
+def test_no_parameter_annotation(x):
+    x.method()
+
+def function_with_return_annotation() -> Foo:
+    return eval("Foo()")
+
+def test_return_annotation():
+    x = function_with_return_annotation() #$ pt,tt=function_with_return_annotation
+    x.method() #$ tt=Foo.method
+
+def function_without_return_annotation():
+    return eval("Foo()")
+
+def test_no_return_annotation():
+    x = function_without_return_annotation() #$ pt,tt=function_without_return_annotation
+    x.method()
+
+def test_variable_annotation():
+    x = eval("Foo()")
+    x : Foo
+    # Currently fails because there is no flow from the class definition to the type annotation.
+    x.method() #$ MISSING: tt=Foo.method
+
+def test_no_variable_annotation():
+    x = eval("Foo()")
+    x.method()