diff --git a/python/ql/lib/semmle/python/Exprs.qll b/python/ql/lib/semmle/python/Exprs.qll index accc370481a..9e00e4f794b 100644 --- a/python/ql/lib/semmle/python/Exprs.qll +++ b/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 | + this = f.getReturns() and result = f.getInnerScope().getReturnNode().getNode() + ) + } } /* Expression Contexts */ diff --git a/python/ql/lib/semmle/python/dataflow/new/internal/DataFlowDispatch.qll b/python/ql/lib/semmle/python/dataflow/new/internal/DataFlowDispatch.qll index 1a38593bce4..781023a9658 100644 --- a/python/ql/lib/semmle/python/dataflow/new/internal/DataFlowDispatch.qll +++ b/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