Python: Post-processing query for inline test expectations

python/ql/test/TestUtilities/InlineExpectationsTest.qll

@@ -5,15 +5,5 @@
 
 private import python as PY
 private import codeql.util.test.InlineExpectationsTest
-
-private module Impl implements InlineExpectationsTestSig {
-  /**
-   * A class representing line comments in Python. As this is the only form of comment Python
-   * permits, we simply reuse the `Comment` class.
-   */
-  class ExpectationComment = PY::Comment;
-
-  class Location = PY::Location;
-}
-
+private import internal.InlineExpectationsTestImpl
 import Make<Impl>

python/ql/test/TestUtilities/InlineExpectationsTestQuery.ql

@@ -0,0 +1,21 @@
+/**
+ * @kind test-postprocess
+ */
+
+private import python
+private import codeql.util.test.InlineExpectationsTest as T
+private import internal.InlineExpectationsTestImpl
+import T::TestPostProcessing
+import T::TestPostProcessing::Make<Impl, Input>
+
+private module Input implements T::TestPostProcessing::InputSig<Impl> {
+  string getRelativeUrl(Location location) {
+    exists(File f, int startline, int startcolumn, int endline, int endcolumn |
+      location.hasLocationInfo(_, startline, startcolumn, endline, endcolumn) and
+      f = location.getFile()
+    |
+      result =
+        f.getRelativePath() + ":" + startline + ":" + startcolumn + ":" + endline + ":" + endcolumn
+    )
+  }
+}

python/ql/test/TestUtilities/internal/InlineExpectationsTestImpl.qll

@@ -0,0 +1,12 @@
+private import python as PY
+private import codeql.util.test.InlineExpectationsTest
+
+module Impl implements InlineExpectationsTestSig {
+  /**
+   * A class representing line comments in Python. As this is the only form of comment Python
+   * permits, we simply reuse the `Comment` class.
+   */
+  class ExpectationComment = PY::Comment;
+
+  class Location = PY::Location;
+}

python/ql/test/query-tests/Numerics/Pythagorean.qlref

@@ -1 +1,2 @@
-Numerics/Pythagorean.ql
+query: Numerics/Pythagorean.ql
+postprocess: TestUtilities/InlineExpectationsTestQuery.ql

python/ql/test/query-tests/Numerics/pythagorean_test.py

@@ -3,12 +3,12 @@
 from math import sqrt
 
 def withPow(a, b):
-    return sqrt(a**2 + b**2)
+    return sqrt(a**2 + b**2) # $ Alert
 
 def withMul(a, b):
-    return sqrt(a*a + b*b)
+    return sqrt(a*a + b*b) # $ Alert
 
 def withRef(a, b):
     a2 = a**2
     b2 = b*b
-    return sqrt(a2 + b2)
+    return sqrt(a2 + b2) # $ Alert

python/ql/test/query-tests/Security/CWE-094-CodeInjection/CodeInjection.qlref

@@ -1 +1,2 @@
-Security/CWE-094/CodeInjection.ql
+query: Security/CWE-094/CodeInjection.ql
+postprocess: TestUtilities/InlineExpectationsTestQuery.ql

python/ql/test/query-tests/Security/CWE-094-CodeInjection/code_injection.py

@@ -1,13 +1,13 @@
-from flask import Flask, request
+from flask import Flask, request # $ Source=flask
 app = Flask(__name__)
 
 @app.route("/code-execution")
 def code_execution():
     code = request.args.get("code")
-    exec(code) # NOT OK
-    eval(code) # NOT OK
+    exec(code) # $ Alert=flask
+    eval(code) # $ Alert=flask
     cmd = compile(code, "<filename>", "exec")
-    exec(cmd) # NOT OK
+    exec(cmd) # $ Alert=flask
 
 
 @app.route("/safe-code-execution")
@@ -18,5 +18,5 @@ def code_execution():
     obj_name = request.args.get("obj")
     if obj_name == "foo" or obj_name == "bar":
         # TODO: Should not alert on this
-        obj = eval(obj_name) # OK
+        obj = eval(obj_name) # $ SPURIOUS: Alert=flask
         print(obj, obj*10)