Help with extracting body of function/methods called in assert statements (python)

[varundataquest]

Using CodeQL, I need help to extract the focal methods (the method that is directly called in assert statements) of the Python unit tests. For example if I have a the following test for the add_list_items function

Test:
def test_add_list_items():
assert add_list_items([1, 2, 3]) == 6

Function:
def add_list_items(z):
return sum(z)

I need to use CodeQL to first extract the method/functions invocations (if there is no method/function, it should skip) inside the assert statements, get the signature of those methods/functions, and then look for them in the project's source code to extract the methods/functions' body as well in CodeQL code.
i.e In the above case, I need to print out
def add_list_items(z):
return sum(z)

Kindly help

[smowton]

How about something like

import python

from Value v, CallNode c, string file, int sl, int el
where c = v.getACall()
and c.getNode() = any(Assert a).getAChildNode*()
and v.hasLocationInfo(file, sl, _, _, _)
and el = max(int line, AstNode a | a = v.getOrigin().getNode().getAChildNode*() and a.getLocation().getEndLine() = line | line)
select c, "Call to $@ at " + file + ":" + sl + "-" + el, v, v.toString()

Note that tells you what functions were called within assert statements, and the file and line ranges those functions occur on, but a CodeQL database does not contain the source code itself. It is therefore necessary for a supporting script to actually extract those lines from the source code.

Note also that the CodeQL database should contain all the relevant application code and of course the tests that contain the assertions, because the full dataflow is used to infer how import statements resolve and therefore the call graph.

[varundataquest]

Thank you very much Chris! This is very helpful.
I have been playing around with this query and I think it needs some more fixes.

• I need to get the signature of those methods/functions as well. The query here does not return that. Is there some changes we can do to get the signature of the methods just with .ql (and not writing external code to get it from source code)
• It returns methods that are not part of the projects (such as Python functions which I am not interested in). How do we limit it to Project source code?
• I agree, I have to write some external code to get to the source and extract the method

An equivalent .ql for Java is something like this:

import java

class AssertMethodCall extends MethodAccess {
AssertMethodCall() {
this.getMethod().getDeclaringType().getQualifiedName().matches("%Assert%")
}
}
from AssertMethodCall assertMethodCall, Method method
where
exists(MethodAccess methodCall |
methodCall = assertMethodCall.getAnArgument().(MethodAccess) and
method = methodCall.getMethod()
)
select
method.getBody().getLocation(),
""+ method.getName()+ "\n"+ method.getSignature() + "\n"+ method.getReturnType() + "\n"+ method.getAModifier()

This does everything needed for Java - Gets the method signature, and also points to where the method is located, using which one can write some code to extract the method out.

[smowton]

Signature: Function has members getArg(int), getKwonlyarg(int) and other related methods (see the definition of class Function). Naturally Python functions don't have a signature in quite the same way as Java functions, but hopefully these will allow you to extract the information you want?

Restricting the particular callees you're interested in: try using Function.getLocation().getFile().getAbsolutePath() (or .getBaseName() if that's sufficient)? You can restrict this as you need to using the string functions matches(string) or regexpMatch(string).

[varundataquest]

Chris,
I took your suggestion and expanded this to check for Function calls that are within the repo and also not an inbuilt function. Also got the Function object by matching with Value by name. I then tried to get the Function name, definition, args and body in the select statement. I ran it using codeQL CLI and I am not getting the function name, definition, args and body in the csv file. I am sure I am doing something wrong here (or that is how this works, I am confused). Kindly help out pl.

/**
* @kind problem
* @id python/quick-query
*/

import python

// Is the Definition of Function in a file within the repo? and is not a python function
// Python functions have a path of python
predicate is_in_repo(Function func) {
    not func.getLocation().getFile().getAbsolutePath().matches("%python%")
  }

// Is it a built-in function?
predicate is_builtin_function(Value v) {
  v.isBuiltin()
}

from Value v, Assert asse, CallNode call, Function func, string file, int sl,  int el
/*
  v is the Value
  a is the assert call, 
  call is the actual call made to in the assert
  func is the function/method definition 
  file is the file info, sl, and el are start and end line of the function/method definition
  Example:

  "a call":
  assert add_list_items([1, 2, 3]) == 6

  "func" :
  def add_list_items(z):
      return sum(z)
  
*/

where
//Get each Function/Method
call = v.getACall()

// Get those where assert is calling the called function/method
and call.getNode() = any(asse).getAChildNode*()

// and that function/method is not a python builtin function like int(), str() len() isInstance() etc.
and not is_builtin_function(v)

// We get the location Info for the called function/method
and v.hasLocationInfo(file, sl, _, _, _)
and el = max(int line, AstNode a | a = v.getOrigin().getNode().getAChildNode*() and a.getLocation().getEndLine() = line | line)


//Get Function that matches the value
and func.getName().matches(v.getName()) // Is this correct way to get a matching function?

// Make sure the called function/method is in the repo and not a call to Python 3.x 

and is_in_repo(func)

select asse, 
       "  $@" + file + ":" + sl + "-" + el + "Name: " + func.getName() + " Def: " + func.getDefinition() + "Args: " + func.getArgs() + " Body: " + func.getBody(),
       func, "Assert"

Example output csv file line:

  |   | warning | [["Assert"\|"relative:///..../tests/integration/support.py:548:5:548:54"]].../tests/integration/support.py:548-561Name: assert_banner_message Def: FunctionExprArgs: ParameterList Body: StmtList | /...tests/integration/test_plugins.py | 370 | 9 | 370 | 55
-- | -- | -- | -- | -- | -- | -- | -- | --

Function is the function that is called by the assert statement. Function Name is printed, Function Definition, Args, Body are not printed. Rest looks ok.

Kindly help out on what I am doing wrong here pl.

[smowton]

To map a Value onto a Function, try using CallableValue instead of Value and CallableValue.getScope(), which returns a Function.

Note that any(asse) doesn't do anything -- just use asse. any means any value that matches the conditions inside the parentheses -- like any(int i | i > 0 and i < 10) for example -- here there are no constraints, so it just evaluates to asse itself.

Re printing things out, the items in question are getting printed, but probably aren't the things you want to print exactly. The Def is printing as FunctionExpr, which means this is an instance of FunctionExpr_, a class without a meaningful pretty-printing in the toString routine. It does have getName, getArgs and other routines to inspect it, but I think this expression probably isn't useful to you -- it's just a lower-level representation of the function definition.

The args are printing as ParameterList, which is class ParameterList -- again, to get information out of this class you need to use the getItem(int) accessor. If you want to glue these together into a string, you might be interested in the concat aggregate described at https://codeql.github.com/docs/ql-language-reference/expressions/#aggregations -- for example, perhaps concat(int i | | v.getScope().(Function).getArg(i).getName(), ", " order by i asc) does what you want?

The same applies to the body -- that's a statement list, which pretty-prints as just StmtList, and has to be interrogated for information about the statements in the list. If you're hoping to render the statement list back out as python source code, I strongly recommend not doing this and instead using an ancillary script to extract the python source code from the original source file based on the source location and the line indices sl and el that are already being produced.

[varundataquest]

Thank you Chris. Your comments really helped a lot.

I did give up on getDefinition() and getBody() on the function as it was not giving useful info like you mentioned. This is the.ql I ended up and on the generated .csv, I was anyway able to write python code to extract out the function given the begin and end line as well as extract out the actual assert call given begin and end line and columns.

I am assuming this .ql code takes care of asserts on independent functions as well as asserts on Function methods (within a class)?

/**
* @kind problem
* @id python/quick-query
*/

import python

// Is the Definition of Function in a file within the repo? and is not a python function
// Python functions have a path of python
predicate is_in_repo(Function func) {
    not func.getLocation().getFile().getAbsolutePath().matches("%python%")
  }

// Is it a built-in function?
predicate is_builtin_function(Value v) {
  v.isBuiltin()
}

from CallableValue v, Assert asse, CallNode call, Function func, string file, int sl,  int el
/*
  v is the CallableValue
  a is the assert call, 
  call is the actual call made to in the assert
  func is the function/method definition 
  file is the file info, sl, and el are start and end line of the function/method definition
  Example:

  "a call":
  assert add_list_items([1, 2, 3]) == 6

  "func" :
  def add_list_items(z):
      return sum(z)
  
*/

where
//Get the CallNode
call = v.getACall()

//Get Function from the Callable Value
and func = v.getScope()

// Get those where assert is calling the called function/method
and call.getNode() = asse.getAChildNode*()

// and that function/method is not a python builtin function like int(), str() len() isInstance() etc.
and not is_builtin_function(v)

// We get the location Info for the called function/method
and v.hasLocationInfo(file, sl, _, _, _)
and el = max(int line, AstNode a | a = v.getOrigin().getNode().getAChildNode*() and a.getLocation().getEndLine() = line | line)


// Make sure the called function/method is in the repo and not a call to Python 3.x 
and is_in_repo(func)


select asse.getLocation(), 
       "  $@ \n" + file + ":" + sl + "-" + el + "\n" + func.getName() + "\n" + concat(int i | | func.getArg(i).getName(), ", " order by i asc),
       func, " Focal"

//Node func.getDefinition() and func.getBody() does not return useful information. We have the extract the function 
//from the file using the fine and sl and el line numbers

[smowton]

I expect so, but I recommend writing up some test cases to make sure it behaves the way you hope.