1

I am trying to traverse an ast of Python source code and check the names of function calls against a whitelist of allowed function names. For example, the following code correctly accomplishes this task:

import ast

# example list of whitelisted functions (real list is much much bigger and includes many libraries)
whitelisted_functions = [
    # From the Math library
    "cos",
    "sin",
    "tan"
]

class Filter(ast.NodeVisitor):
    def visit_Call(self, node):
        if node.func.id not in whitelisted_functions:
            # throw some exception
        self.generic_visit(node)

def filter_code(code):
    tree = ast.parse(code)

    filter = Filter()
    filter.visit(tree)

However, I further want to get the name of the module which the function definition is in. This is because when the code checks if the function ID is in the list of whitelisted functions, I further want to check that the function definition is from a specific module. For example, when the code checks that "cos" is in the list of whitelisted functions, I want to get the name of the module containing the function definition and verify that it is in fact "math".

I have been perusing StackOverflow and various documentation online (e.g. https://docs.python.org/dev/library/ast.html); however, I am still unsure how to accomplish this task. Is there any way to do this? Any help would be greatly appreciated!

RGV
  • 11
  • 4
  • That's a strange question. A function's signature isn't defined somewhere separate from the function. Maybe you have a misunderstanding of what the word "signature" means. – user2357112 Jul 14 '20 at 21:55
  • @user2357112supportsMonica You're absolutely correct -- my bad. I meant the function definition, not the signature -- I'll edit the question to reflect this. – RGV Jul 14 '20 at 21:57
  • An AST `Call` node for something like `math.cos(...)` isn't directly going to have a `func.id` attribute that's of any use. Its `func` will be an `Attribute` node, whose `value` is a `Name` with `id` "math", and whose `attr` is "cos". – jasonharper Jul 14 '20 at 22:10

0 Answers0