Python: Handle defaults and annotations for keyword-only arguments

This commit is based on a change to the extractor

Author: RasmusWL

python/ql/src/semmle/python/AstGenerated.qll

@@ -1191,13 +1191,13 @@ library class AliasList_ extends @py_alias_list {
 }
 
 library class Arguments_ extends @py_arguments {
-    /** Gets the keyword default values of this parameters definition. */
+    /** Gets the keyword-only default values of this parameters definition. */
     ExprList getKwDefaults() { py_expr_lists(result, this, 0) }
 
-    /** Gets the nth keyword default value of this parameters definition. */
+    /** Gets the nth keyword-only default value of this parameters definition. */
     Expr getKwDefault(int index) { result = this.getKwDefaults().getItem(index) }
 
-    /** Gets a keyword default value of this parameters definition. */
+    /** Gets a keyword-only default value of this parameters definition. */
     Expr getAKwDefault() { result = this.getKwDefaults().getAnItem() }
 
     /** Gets the default values of this parameters definition. */
@@ -1224,13 +1224,13 @@ library class Arguments_ extends @py_arguments {
     /** Gets the **kwarg annotation of this parameters definition. */
     Expr getKwargannotation() { py_exprs(result, _, this, 4) }
 
-    /** Gets the kw_annotations of this parameters definition. */
+    /** Gets the keyword-only annotations of this parameters definition. */
     ExprList getKwAnnotations() { py_expr_lists(result, this, 5) }
 
-    /** Gets the nth kw_annotation of this parameters definition. */
+    /** Gets the nth keyword-only annotation of this parameters definition. */
     Expr getKwAnnotation(int index) { result = this.getKwAnnotations().getItem(index) }
 
-    /** Gets a kw_annotation of this parameters definition. */
+    /** Gets a keyword-only annotation of this parameters definition. */
     Expr getAKwAnnotation() { result = this.getKwAnnotations().getAnItem() }
 
     ArgumentsParent getParent() { py_arguments(this, result) }

python/ql/src/semmle/python/Function.qll

@@ -209,19 +209,31 @@ class Parameter extends Parameter_ {
 
     /** Gets the expression for the default value of this parameter */
     Expr getDefault() {
-        exists(Function f, int n, int c, int d, Arguments args | args = f.getDefinition().getArgs() |
-            f.getArg(n) = this and
-            c = count(f.getAnArg()) and
-            d = count(args.getADefault()) and
-            result = args.getDefault(d - c + n)
+        exists(Function f, int i, Arguments args | args = f.getDefinition().getArgs() |
+            // positional (normal)
+            f.getArg(i) = this and
+            result = args.getDefault(i)
+        )
+        or
+        exists(Function f, int i, Arguments args | args = f.getDefinition().getArgs() |
+            // keyword-only
+            f.getKeywordOnlyArg(i) = this and
+            result = args.getKwDefault(i)
         )
     }
 
     /** Gets the annotation expression of this parameter */
     Expr getAnnotation() {
-        exists(Function f, int n, Arguments args | args = f.getDefinition().getArgs() |
-            f.getArg(n) = this and
-            result = args.getAnnotation(n)
+        exists(Function f, int i, Arguments args | args = f.getDefinition().getArgs() |
+            // positional (normal)
+            f.getArg(i) = this and
+            result = args.getAnnotation(i)
+        )
+        or
+        exists(Function f, int i, Arguments args | args = f.getDefinition().getArgs() |
+            // keyword-only
+            f.getKeywordOnlyArg(i) = this and
+            result = args.getKwAnnotation(i)
         )
         or
         exists(Function f, Arguments args | args = f.getDefinition().getArgs() |
@@ -342,4 +354,19 @@ class Arguments extends Arguments_ {
         result = this.getAKwAnnotation() or
         result = this.getKwargannotation()
     }
+
+    // The following 4 methods are overwritten to provide better QLdoc. Since the
+    // Arguments_ is auto-generated, we can't change the poor auto-generated docs there :(
+
+    /** Gets the default value for the `index`'th positional parameter. */
+    override Expr getDefault(int index) { result = super.getDefault(index) }
+
+    /** Gets the default value for the `index`'th keyword-only parameter. */
+    override Expr getKwDefault(int index) { result = super.getKwDefault(index) }
+
+    /** Gets the annotation for the `index`'th positional parameter. */
+    override Expr getAnnotation(int index) { result = super.getAnnotation(index) }
+
+    /** Gets the annotation for the `index`'th keyword-only parameter. */
+    override Expr getKwAnnotation(int index) { result = super.getKwAnnotation(index) }
 }

python/ql/test/3/library-tests/functions/FunctionExpr.getArgs.getDefault.expected

@@ -1,4 +1,4 @@
 | test.py:4:1:11:2 | FunctionExpr | 0 | test.py:5:21:5:22 | UnaryExpr |
 | test.py:4:1:11:2 | FunctionExpr | 1 | test.py:7:19:7:20 | UnaryExpr |
-| test.py:23:1:31:2 | FunctionExpr | 0 | test.py:25:28:25:31 | None |
-| test.py:23:1:31:2 | FunctionExpr | 1 | test.py:26:20:26:23 | None |
+| test.py:23:1:31:2 | FunctionExpr | 1 | test.py:25:28:25:31 | None |
+| test.py:23:1:31:2 | FunctionExpr | 2 | test.py:26:20:26:23 | None |

python/ql/test/3/library-tests/functions/FunctionExpr.getArgs.getKwDefault.expected

@@ -1,2 +1,2 @@
 | test.py:4:1:11:2 | FunctionExpr | 0 | test.py:9:25:9:26 | UnaryExpr |
-| test.py:23:1:31:2 | FunctionExpr | 0 | test.py:29:32:29:35 | None |
+| test.py:23:1:31:2 | FunctionExpr | 1 | test.py:29:32:29:35 | None |

python/ql/test/3/library-tests/parameters/Annotations.expected

@@ -1,4 +1,6 @@
 | args | test.py:8:12:8:23 | Str |
+| keyword_only | test.py:9:19:9:21 | int |
+| keyword_w_default | test.py:29:24:29:28 | Str |
 | kwargs | test.py:10:15:10:30 | Str |
 | normal | test.py:7:13:7:15 | int |
 | pos_only | test.py:5:15:5:17 | int |

python/ql/test/3/library-tests/parameters/Defaults.expected

@@ -1,3 +1,5 @@
+| keyword_only | test.py:9:25:9:26 | UnaryExpr |
+| keyword_w_default | test.py:29:32:29:35 | None |
 | normal | test.py:7:19:7:20 | UnaryExpr |
 | pos_only | test.py:5:21:5:22 | UnaryExpr |
 | pos_w_default | test.py:25:28:25:31 | None |