view src/ast_convert.rs @ 9:fa7e285f88e7

Add a scoping pass, associating each module/statement with a block.
author Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
date Wed, 01 Jun 2016 22:17:28 +0100
parents 94ff501bf336
children 38b0d63697b1
line wrap: on
line source

use python_ast::{Module, Statement, Expr, BinOp, UnaryOp};

use cpython::{Python, PyObject};
use cpython::ObjectProtocol; //for call method

fn get_str(py: Python, object: PyObject) -> String {
    let pystring = object.str(py).unwrap();
    let mut string = pystring.to_string(py).unwrap();
    string.to_mut().to_string()
}

fn parse_expr_vec(py: Python, ast: PyObject) -> Vec<String> {
    let builtins_module = py.import("builtins").unwrap();
    let isinstance = builtins_module.get(py, "isinstance").unwrap();

    let is_instance = |object: &PyObject, type_: &PyObject| {
        return isinstance.call(py, (object, type_), None).unwrap().is_true(py).unwrap();
    };

    let ast_module = py.import("ast").unwrap();
    let ast_type = ast_module.get(py, "AST").unwrap();
    let arguments_type = ast_module.get(py, "arguments").unwrap();

    assert!(is_instance(&ast, &ast_type));

    if is_instance(&ast, &arguments_type) {
        let args = ast.getattr(py, "args").unwrap();
        let mut arguments = vec!();
        for arg in args.iter(py).unwrap() {
            let arg = arg.unwrap();
            let arg = parse_expr(py, arg);
            arguments.push(match arg {
                Expr::Name(arg) => arg,
                _ => panic!()
            });
        }
        arguments
    } else {
        panic!()
    }
}

fn parse_unaryop(py: Python, ast: PyObject) -> UnaryOp {
    let builtins_module = py.import("builtins").unwrap();
    let isinstance = builtins_module.get(py, "isinstance").unwrap();

    let is_instance = |object: &PyObject, type_: &PyObject| {
        return isinstance.call(py, (object, type_), None).unwrap().is_true(py).unwrap();
    };

    let ast_module = py.import("ast").unwrap();
    let ast_type = ast_module.get(py, "AST").unwrap();
    let uadd_type = ast_module.get(py, "UAdd").unwrap();
    let usub_type = ast_module.get(py, "USub").unwrap();

    assert!(is_instance(&ast, &ast_type));

    if is_instance(&ast, &uadd_type) {
        UnaryOp::UAdd
    } else if is_instance(&ast, &usub_type) {
        UnaryOp::USub
    } else {
        println!("UnaryOp {}", ast);
        UnaryOp::Error
    }
}

fn parse_binop(py: Python, ast: PyObject) -> BinOp {
    let builtins_module = py.import("builtins").unwrap();
    let isinstance = builtins_module.get(py, "isinstance").unwrap();

    let is_instance = |object: &PyObject, type_: &PyObject| {
        return isinstance.call(py, (object, type_), None).unwrap().is_true(py).unwrap();
    };

    let ast_module = py.import("ast").unwrap();
    let ast_type = ast_module.get(py, "AST").unwrap();
    let add_type = ast_module.get(py, "Add").unwrap();
    let mult_type = ast_module.get(py, "Mult").unwrap();
    let eq_type = ast_module.get(py, "Eq").unwrap();
    let lt_type = ast_module.get(py, "Lt").unwrap();
    let gt_type = ast_module.get(py, "Gt").unwrap();
    let sub_type = ast_module.get(py, "Sub").unwrap();
    let div_type = ast_module.get(py, "Div").unwrap();

    assert!(is_instance(&ast, &ast_type));

    if is_instance(&ast, &add_type) {
        BinOp::BinAdd
    } else if is_instance(&ast, &mult_type) {
        BinOp::BinMult
    } else if is_instance(&ast, &eq_type) {
        BinOp::BinEq
    } else if is_instance(&ast, &lt_type) {
        BinOp::BinLt
    } else if is_instance(&ast, &gt_type) {
        BinOp::BinGt
    } else if is_instance(&ast, &sub_type) {
        BinOp::Sub
    } else if is_instance(&ast, &div_type) {
        BinOp::Div
    } else {
        println!("BinOp {}", ast);
        BinOp::Error
    }
}

fn parse_expr(py: Python, ast: PyObject) -> Expr {
    let builtins_module = py.import("builtins").unwrap();
    let isinstance = builtins_module.get(py, "isinstance").unwrap();

    let is_instance = |object: &PyObject, type_: &PyObject| {
        return isinstance.call(py, (object, type_), None).unwrap().is_true(py).unwrap();
    };

    let ast_module = py.import("ast").unwrap();
    let ast_type = ast_module.get(py, "AST").unwrap();
    let arg_type = ast_module.get(py, "arg").unwrap();
    let unary_op_type = ast_module.get(py, "UnaryOp").unwrap();
    let bin_op_type = ast_module.get(py, "BinOp").unwrap();
    let name_constant_type = ast_module.get(py, "NameConstant").unwrap();
    let attribute_type = ast_module.get(py, "Attribute").unwrap();
    let name_type = ast_module.get(py, "Name").unwrap();
    let num_type = ast_module.get(py, "Num").unwrap();
    let str_type = ast_module.get(py, "Str").unwrap();
    let list_type = ast_module.get(py, "List").unwrap();
    let compare_type = ast_module.get(py, "Compare").unwrap();
    let call_type = ast_module.get(py, "Call").unwrap();
    let alias_type = ast_module.get(py, "alias").unwrap();

    assert!(is_instance(&ast, &ast_type));

    if is_instance(&ast, &arg_type) {
        let arg = ast.getattr(py, "arg").unwrap();
        let arg = get_str(py, arg);
        Expr::Name(arg)
    } else if is_instance(&ast, &attribute_type) {
        let value = ast.getattr(py, "value").unwrap();
        let attr = ast.getattr(py, "attr").unwrap();

        let value = parse_expr(py, value);
        let attr = get_str(py, attr);

        Expr::Attribute(Box::new(value), attr)
    } else if is_instance(&ast, &name_type) {
        let id = ast.getattr(py, "id").unwrap();
        let id = get_str(py, id);
        Expr::Name(id)
    } else if is_instance(&ast, &name_constant_type) {
        let value = ast.getattr(py, "value").unwrap();
        let value = get_str(py, value);
        Expr::NameConstant(value)
    } else if is_instance(&ast, &num_type) {
        let n = ast.getattr(py, "n").unwrap();
        let n = get_str(py, n);
        Expr::Num(n)
    } else if is_instance(&ast, &str_type) {
        let s = ast.getattr(py, "s").unwrap();
        let s = get_str(py, s);
        Expr::Str(s)
    } else if is_instance(&ast, &list_type) {
        let elts = ast.getattr(py, "elts").unwrap();

        let mut elements = vec!();
        for elt in elts.iter(py).unwrap() {
            let elt = elt.unwrap();
            elements.push(parse_expr(py, elt));
        }

        Expr::List(elements)
    } else if is_instance(&ast, &unary_op_type) {
        let op = ast.getattr(py, "op").unwrap();
        let operand = ast.getattr(py, "operand").unwrap();

        let op = parse_unaryop(py, op);
        let operand = parse_expr(py, operand);

        Expr::UnaryOp(op, Box::new(operand))
    } else if is_instance(&ast, &bin_op_type) {
        let left = ast.getattr(py, "left").unwrap();
        let op = ast.getattr(py, "op").unwrap();
        let right = ast.getattr(py, "right").unwrap();

        let left = parse_expr(py, left);
        let op = parse_binop(py, op);
        let right = parse_expr(py, right);

        Expr::BinOp(Box::new(left), op, Box::new(right))
    } else if is_instance(&ast, &call_type) {
        let func = ast.getattr(py, "func").unwrap();
        let args = ast.getattr(py, "args").unwrap();
        //let keywords = ast.getattr(py, "keywords").unwrap();

        let func = parse_expr(py, func);

        let mut arguments = vec!();
        for arg in args.iter(py).unwrap() {
            let arg = arg.unwrap();
            arguments.push(parse_expr(py, arg));
        }

        Expr::Call(Box::new(func), arguments)
    } else if is_instance(&ast, &alias_type) {
        let name = ast.getattr(py, "name").unwrap();
        let asname = ast.getattr(py, "asname").unwrap();

        let name = get_str(py, name);
        let asname = if asname == py.None() {
            None
        } else {
            Some(get_str(py, asname))
        };

        Expr::Alias(name, asname)
    } else if is_instance(&ast, &compare_type) {
        let left = ast.getattr(py, "left").unwrap();
        let ops = ast.getattr(py, "ops").unwrap();
        let comparators = ast.getattr(py, "comparators").unwrap();

        let left = parse_expr(py, left);
        let ops = ops.iter(py).unwrap();
        let comparators = comparators.iter(py).unwrap();

        let mut new_ops = vec!();
        for op in ops {
            let op = op.unwrap();
            let op = parse_binop(py, op);
            new_ops.push(op);
        }

        let mut new_comparators = vec!();
        for comparator in comparators {
            let comparator = comparator.unwrap();
            let comparator = parse_expr(py, comparator);
            new_comparators.push(comparator);
        }

        Expr::Compare(Box::new(left), new_ops, new_comparators)
    } else {
        println!("Expr {}", ast);
        Expr::Error
    }
}

fn parse_statement(py: Python, ast: PyObject) -> Statement {
    //Statement::FunctionDef(Expr::Name("function".to_string()), vec!(Expr::Name("a".to_string()), Expr::Name("b".to_string())), vec!())
    //Statement::If(Expr::BinOp(BinOp::BinEq, Box::new(Expr::Name("__name__".to_string())), Box::new(Expr::Str("__main__".to_string()))), vec!(Statement::Expr(Expr::Call(Box::new(Expr::Name("function".to_string())), vec!(Expr::Num(1), Expr::Num(2))))))

    let builtins_module = py.import("builtins").unwrap();
    let isinstance = builtins_module.get(py, "isinstance").unwrap();

    let is_instance = |object: &PyObject, type_: &PyObject| {
        return isinstance.call(py, (object, type_), None).unwrap().is_true(py).unwrap();
    };

    let ast_module = py.import("ast").unwrap();
    let ast_type = ast_module.get(py, "AST").unwrap();
    let class_def_type = ast_module.get(py, "ClassDef").unwrap();
    let function_def_type = ast_module.get(py, "FunctionDef").unwrap();
    let global_type = ast_module.get(py, "Global").unwrap();
    let assign_type = ast_module.get(py, "Assign").unwrap();
    let aug_assign_type = ast_module.get(py, "AugAssign").unwrap();
    let return_type = ast_module.get(py, "Return").unwrap();
    let import_from_type = ast_module.get(py, "ImportFrom").unwrap();
    let if_type = ast_module.get(py, "If").unwrap();
    let while_type = ast_module.get(py, "While").unwrap();
    let for_type = ast_module.get(py, "For").unwrap();
    let expr_type = ast_module.get(py, "Expr").unwrap();
    let break_type = ast_module.get(py, "Break").unwrap();

    assert!(is_instance(&ast, &ast_type));

    /*
    // TODO: implement Hash for PyObject. (trivial)
    let map = {
        let fields = ast.getattr(py, "_fields").unwrap();
        let mut map = HashMap::new();
        for field in fields.iter(py).unwrap() {
            let field = field.unwrap();
            let value = ast.getattr(py, field).unwrap();
            map.insert(field, value);
        }
        map
    };
    */

    if is_instance(&ast, &class_def_type) {
        let name = ast.getattr(py, "name").unwrap();
        let bases = ast.getattr(py, "bases").unwrap();
        //let keywords = ast.getattr(py, "keywords").unwrap();
        let body = ast.getattr(py, "body").unwrap();
        //let decorator_list = ast.getattr(py, "decorator_list").unwrap();

        let name = get_str(py, name);

        let mut nodes = vec!();
        for name_node in bases.iter(py).unwrap() {
            let name_node = name_node.unwrap();
            let name_node = parse_expr(py, name_node);
            nodes.push(name_node);
        }

        let mut statements = vec!();
        for statement in body.iter(py).unwrap() {
            let statement = statement.unwrap();
            let statement = parse_statement(py, statement);
            statements.push(statement);
        }

        Statement::ClassDef(name, nodes, statements)
    } else if is_instance(&ast, &function_def_type) {
        let name = ast.getattr(py, "name").unwrap();
        let args = ast.getattr(py, "args").unwrap();
        let body = ast.getattr(py, "body").unwrap();

        let name = get_str(py, name);
        let args = parse_expr_vec(py, args);
        /*
        let mut arguments = vec!();
        for arg in args.iter(py).unwrap() {
            let arg = parse_expr(py, arg.unwrap());
            arguments.push(arg);
        }
        */

        let mut statements = vec!();
        for statement in body.iter(py).unwrap() {
            let statement = parse_statement(py, statement.unwrap());
            statements.push(statement);
        }

        Statement::FunctionDef(name, args, statements)
    } else if is_instance(&ast, &global_type) {
        let names = ast.getattr(py, "names").unwrap();

        let mut globals = vec!();
        for name in names.iter(py).unwrap() {
            let name = name.unwrap();
            let name = get_str(py, name);
            globals.push(name);
        }

        Statement::Global(globals)
    } else if is_instance(&ast, &if_type) {
        let test = ast.getattr(py, "test").unwrap();
        let body = ast.getattr(py, "body").unwrap();
        let orelse = ast.getattr(py, "orelse").unwrap();

        let test = parse_expr(py, test);

        let mut statements = vec!();
        for statement in body.iter(py).unwrap() {
            let statement = statement.unwrap();
            let statement = parse_statement(py, statement);
            statements.push(statement);
        }

        let mut orelse_ = vec!();
        for statement in orelse.iter(py).unwrap() {
            let statement = statement.unwrap();
            let statement = parse_statement(py, statement);
            orelse_.push(statement);
        }

        Statement::If(test, statements, orelse_)
    } else if is_instance(&ast, &while_type) {
        let test = ast.getattr(py, "test").unwrap();
        let body = ast.getattr(py, "body").unwrap();
        let orelse = ast.getattr(py, "orelse").unwrap();

        let test = parse_expr(py, test);

        let mut statements = vec!();
        for statement in body.iter(py).unwrap() {
            let statement = statement.unwrap();
            let statement = parse_statement(py, statement);
            statements.push(statement);
        }

        let mut orelse_ = vec!();
        for statement in orelse.iter(py).unwrap() {
            let statement = statement.unwrap();
            let statement = parse_statement(py, statement);
            orelse_.push(statement);
        }

        Statement::While(test, statements, orelse_)
    } else if is_instance(&ast, &for_type) {
        let target = ast.getattr(py, "target").unwrap();
        let iter = ast.getattr(py, "iter").unwrap();
        let body = ast.getattr(py, "body").unwrap();
        let orelse = ast.getattr(py, "orelse").unwrap();

        let target = parse_expr(py, target);
        let iter = parse_expr(py, iter);

        let mut statements = vec!();
        for statement in body.iter(py).unwrap() {
            let statement = statement.unwrap();
            let statement = parse_statement(py, statement);
            statements.push(statement);
        }

        let mut orelse_ = vec!();
        for statement in orelse.iter(py).unwrap() {
            let statement = statement.unwrap();
            let statement = parse_statement(py, statement);
            orelse_.push(statement);
        }

        Statement::For(target, iter, statements, orelse_)
    } else if is_instance(&ast, &assign_type) {
        let targets = ast.getattr(py, "targets").unwrap();
        let value = ast.getattr(py, "value").unwrap();

        let mut arguments = vec!();
        for target in targets.iter(py).unwrap() {
            let target = target.unwrap();
            let target = parse_expr(py, target);
            arguments.push(target);
        }

        let value = parse_expr(py, value);

        Statement::Assign(arguments, value)
    } else if is_instance(&ast, &aug_assign_type) {
        let target = ast.getattr(py, "target").unwrap();
        let op = ast.getattr(py, "op").unwrap();
        let value = ast.getattr(py, "value").unwrap();

        let target = parse_expr(py, target);
        let op = parse_binop(py, op);
        let value = parse_expr(py, value);

        Statement::AugAssign(target, op, value)
    } else if is_instance(&ast, &import_from_type) {
        let module = ast.getattr(py, "module").unwrap();
        let names = ast.getattr(py, "names").unwrap();
        //let level = ast.getattr(py, "level").unwrap();

        let module = get_str(py, module);

        let mut names_ = vec!();
        for alias in names.iter(py).unwrap() {
            let alias = alias.unwrap();
            let alias = parse_expr(py, alias);
            names_.push(alias);
        }

        Statement::ImportFrom(module, names_)
    } else if is_instance(&ast, &return_type) {
        let value = ast.getattr(py, "value").unwrap();
        let value = parse_expr(py, value);
        Statement::Return(value)
    } else if is_instance(&ast, &expr_type) {
        let value = ast.getattr(py, "value").unwrap();
        let value = parse_expr(py, value);
        Statement::Expr(value)
    } else if is_instance(&ast, &break_type) {
        Statement::Break
    } else {
        println!("Statement {}", ast);
        Statement::Error
    }
}

#[allow(dead_code)]
pub fn convert_ast(name: String, module: &PyObject) -> Module {
    let gil = Python::acquire_gil();
    let py = gil.python();

    let builtins_module = py.import("builtins").unwrap();
    let isinstance = builtins_module.get(py, "isinstance").unwrap();

    let ast_module = py.import("ast").unwrap();
    let module_type = ast_module.get(py, "Module").unwrap();

    assert!(isinstance.call(py, (module, module_type), None).unwrap().is_true(py).unwrap());

    let body = module.getattr(py, "body").unwrap();
    let mut statements = vec!();
    for statement in body.iter(py).unwrap() {
        let statement = parse_statement(py, statement.unwrap());
        statements.push(statement)
    }
    Module{name: name, statements: statements}
}