Python: Better API for parameters.

Author: markshannon

python/ql/src/semmle/python/AstExtended.qll

@@ -81,7 +81,7 @@ class Comprehension extends Comprehension_, AstNode {
 
 }
 
-class BytesOrStr extends BytesOrStr_ { 
+class BytesOrStr extends BytesOrStr_ {
 
 }
 

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

@@ -224,6 +224,12 @@ class Parameter extends Parameter_ {
     Parameter() {
         /* Parameter_ is just defined as a Name or Tuple, narrow to actual parameters */
         exists(ParameterList pl | py_exprs(this, _, pl, _))
+        or
+        exists(Function f |
+            f.getVararg() = this
+            or
+            f.getKwarg() = this
+        )
     }
 
     Location getLocation() {
@@ -242,6 +248,7 @@ class Parameter extends Parameter_ {
         result = this
     }
 
+    /** 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() |
@@ -252,6 +259,24 @@ class Parameter extends Parameter_ {
         )
     }
 
+    /** 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)
+        )
+        or
+        exists(Function f, Arguments args |
+            args = f.getDefinition().getArgs() |
+            f.getKwarg() = this and
+            result = args.getKwargannotation()
+            or
+            f.getVararg() = this and
+            result = args.getVarargannotation()
+        )
+    }
+
     Variable getVariable() {
         result.getAnAccess() = this.asName()
     }

python/ql/src/semmle/python/pointsto/Base.qll

@@ -275,14 +275,19 @@ class ParameterDefinition extends PyNodeDefinition {
         this.getDefiningNode().getNode().(Parameter).isSelf()
     }
 
+    /** Gets the control flow node for the default value of this parameter */
     ControlFlowNode getDefault() {
-        exists(Function f, int n, int c, int d, Arguments args |
-            args = f.getDefinition().getArgs() |
-            f.getArg(n) = this.getDefiningNode().getNode() and
-            c = count(f.getAnArg()) and
-            d = count(args.getADefault()) and
-            result.getNode() = args.getDefault(d-c+n)
-        )
+        result.getNode() = this.getParameter().getDefault()
+    }
+
+    /** Gets the annotation control flow node of this parameter */
+    ControlFlowNode getAnnotation() {
+        result.getNode() = this.getParameter().getAnnotation()
+    }
+
+    /** Gets the name of this parameter definition */
+    string getName() {
+        result = this.getParameter().asName().getId()
     }
 
     predicate isVarargs() {

python/ql/test/library-tests/PointsTo/new/SourceNodeDefinitions.expected

@@ -18,9 +18,7 @@
 | a_simple.py:11 | Global Variable C | ControlFlowNode for C | definition |
 | a_simple.py:14 | Global Variable vararg_kwarg | ControlFlowNode for vararg_kwarg | definition |
 | a_simple.py:14 | Local Variable d | ControlFlowNode for d | definition |
-| a_simple.py:14 | Local Variable d | Entry node for Function vararg_kwarg | definition |
 | a_simple.py:14 | Local Variable t | ControlFlowNode for t | definition |
-| a_simple.py:14 | Local Variable t | Entry node for Function vararg_kwarg | definition |
 | a_simple.py:18 | Global Variable multi_loop | ControlFlowNode for multi_loop | definition |
 | a_simple.py:18 | Local Variable seq | ControlFlowNode for seq | definition |
 | a_simple.py:18 | Local Variable x | Entry node for Function multi_loop | definition |
@@ -40,9 +38,7 @@
 | a_simple.py:29 | Local Variable q | ControlFlowNode for q | definition |
 | a_simple.py:34 | Global Variable f | ControlFlowNode for f | definition |
 | a_simple.py:34 | Local Variable args | ControlFlowNode for args | definition |
-| a_simple.py:34 | Local Variable args | Entry node for Function f | definition |
 | a_simple.py:34 | Local Variable kwargs | ControlFlowNode for kwargs | definition |
-| a_simple.py:34 | Local Variable kwargs | Entry node for Function f | definition |
 | a_simple.py:38 | Global Variable multi_assign_and_packing | ControlFlowNode for multi_assign_and_packing | definition |
 | a_simple.py:38 | Local Variable a | ControlFlowNode for a | definition |
 | a_simple.py:38 | Local Variable b | ControlFlowNode for b | definition |
@@ -142,7 +138,6 @@
 | b_condition.py:99 | Local Variable x | ControlFlowNode for use() | refinement |
 | b_condition.py:101 | Global Variable not_or_not | ControlFlowNode for not_or_not | definition |
 | b_condition.py:101 | Local Variable a | ControlFlowNode for a | definition |
-| b_condition.py:101 | Local Variable a | Entry node for Function not_or_not | definition |
 | b_condition.py:109 | Global Variable method_check | ControlFlowNode for method_check | definition |
 | b_condition.py:109 | Local Variable x | ControlFlowNode for x | definition |
 | b_condition.py:111 | Local Variable x | ControlFlowNode for use() | refinement |

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

@@ -0,0 +1,4 @@
+| args | test.py:5:13:5:22 | Subscript |
+| b | test.py:2:12:2:14 | str |
+| d | test.py:2:27:2:29 | int |
+| kwargs | test.py:5:34:5:47 | Subscript |

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

@@ -0,0 +1,6 @@
+
+import python
+
+
+from Parameter p
+select p.getName(), p.getAnnotation()

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

@@ -0,0 +1,2 @@
+| c | test.py:2:19:2:22 | None |
+| d | test.py:2:31:2:34 | True |

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

@@ -0,0 +1,6 @@
+
+import python
+
+
+from Parameter p
+select p.getName(), p.getDefault()

python/ql/test/library-tests/parameters/Special.expected

@@ -0,0 +1,6 @@
+| a | normal |
+| args | varargs |
+| b | normal |
+| c | normal |
+| d | normal |
+| kwargs | kwargs |

python/ql/test/library-tests/parameters/Special.ql

@@ -0,0 +1,13 @@
+
+
+import python
+
+from Parameter p, string type
+where
+p.isKwargs() and type = "kwargs"
+or
+p.isVarargs() and type = "varargs"
+or
+not p.isKwargs() and not p.isVarargs() and type = "normal"
+
+select p.getName(), type