preprocess.py 文件源码

def preprocess_input_python(s):
    rootnode = py_ast.get_ast(s)
    rootnode = add_str_node(rootnode)

    for node in py_ast.find_all(rootnode, ast.BinOp):
        if isinstance(node.op, ast.Pow):
            if not py_ast.is_int_constant_py_ast(py_ast.dump_ast(node.right)):
                new_str = 'pow(' + py_ast.dump_ast(node.left) + ', ' + py_ast.dump_ast(node.right) + ')'
                new_node = py_ast.get_ast(new_str).body[0].value
                py_ast.replace_node(rootnode, node, new_node)

    for node in py_ast.find_all(rootnode, ast.Call):
        if isinstance(node.func, ast.Attribute):
            if isinstance(node.func.value, ast.Name) and node.func.attr in ['floow', 'ceil']:
                if node.func.value.id == 'math':
                    py_ast.replace_node(node, node.func.value, py_ast.get_ast('numpy').body[0].value)

    ans = py_ast.dump_ast(rootnode)
    ans = 'import numpy\n' + ans

    lines = ans.split('\n')
    i = 0
    while i < len(lines):
        line = lines[i]
        if line.startswith('#'):
            j = i + 1
            while j < len(lines) and lines[j].strip() == '':
                del lines[j]
        i += 1

    ans = '\n'.join(lines)

    return ans