get_calls_from_expr_or_assign

mail@pastecode.io avatar
unknown
python
2 years ago
2.1 kB
4
Indexable
Never

def get_calls_from_expr_or_assign(expr):
    try:
        if type(expr) is ast.Name:
            return [expr.id]
        if type(expr) is ast.Call:
            # function call
            if type(expr.func) is ast.Name:
                name = [expr.func.id]
            # method
            elif type(expr.func) is ast.Attribute:
                name = [expr.func.attr]
            else:
                name = []
            return name + get_calls_from_expr_or_assign(expr.args)
        elif type(expr) is ast.Attribute:
            return [expr.attr]
        elif type(expr) is ast.BinOp:
            return get_calls_from_expr_or_assign(
                expr.left
            ) + get_calls_from_expr_or_assign(expr.right)
        elif type(expr) is ast.Expr:
            return get_calls_from_expr_or_assign(expr.value)
        elif type(expr) is ast.Assign:
            return get_calls_from_expr_or_assign(expr.value)
        elif type(expr) is ast.While or type(expr) is ast.If:
            return (
                get_calls_from_expr_or_assign(expr.test)
                + get_calls_from_expr_or_assign(expr.orelse)
                + get_calls_from_expr_or_assign(expr.body)
            )
        elif type(expr) is ast.For:
            return (
                get_calls_from_expr_or_assign(expr.target)
                + get_calls_from_expr_or_assign(expr.body)
                + get_calls_from_expr_or_assign(expr.iter)
            )
        elif type(expr) is ast.Try:
            return get_calls_from_expr_or_assign(expr.body)
        elif type(expr) is ast.If:
            return get_calls_from_expr_or_assign(
                expr.body
            ) + get_calls_from_expr_or_assign(expr.orelse)
        elif type(expr) is list:
            return [
                res
                for subexpr in expr
                for res in get_calls_from_expr_or_assign(subexpr)
            ]
        else:
            return []
    except Exception as e:
        logging.info(str(astunparse.unparse(expr)), " ", str(e))
        return []