view src/python_dump.rs @ 69:a73eaf42bea1

Add ast.AsyncDefFunction.
author Bastien Orivel <eijebong@bananium.fr>
date Mon, 13 Jun 2016 17:54:19 +0200
parents a0b23123901b
children
line wrap: on
line source

extern crate cpython;

use std::iter;

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

fn dump(py: Python, indent: usize, ast: PyObject) -> 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 function_def_type = ast_module.get(py, "FunctionDef").unwrap();
    let arguments_type = ast_module.get(py, "arguments").unwrap();
    let arg_type = ast_module.get(py, "arg").unwrap();
    let assign_type = ast_module.get(py, "Assign").unwrap();
    let return_type = ast_module.get(py, "Return").unwrap();
    let bin_op_type = ast_module.get(py, "BinOp").unwrap();
    let name_constant_type = ast_module.get(py, "NameConstant").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 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 if_type = ast_module.get(py, "If").unwrap();
    let for_type = ast_module.get(py, "For").unwrap();
    let compare_type = ast_module.get(py, "Compare").unwrap();
    let expr_type = ast_module.get(py, "Expr").unwrap();
    let call_type = ast_module.get(py, "Call").unwrap();
    let import_from_type = ast_module.get(py, "ImportFrom").unwrap();
    let alias_type = ast_module.get(py, "alias").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
    };
    */

    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()
    }

    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 args = dump(py, indent, args);
        let declaration = format!("def {}({}):", name, args);
        let mut statements = vec!(String::new());
        for statement in body.iter(py).unwrap() {
            let statement = dump(py, indent + 1, statement.unwrap());
            statements.push(statement);
        }
        let indent: String = iter::repeat("    ").take(indent + 1).collect();
        let indent = format!("\n{}", indent);
        let indent = indent.as_str();
        let body = statements.join(indent);
        format!("{}{}", declaration, body)
    } else if is_instance(&ast, &if_type) {
        let test = ast.getattr(py, "test").unwrap();
        let body = ast.getattr(py, "body").unwrap();

        let test = dump(py, indent, test);

        let declaration = format!("if {}:", test);
        let mut statements = vec!(String::new());
        for statement in body.iter(py).unwrap() {
            let statement = dump(py, indent + 1, statement.unwrap());
            statements.push(statement);
        }
        let indent: String = iter::repeat("    ").take(indent + 1).collect();
        let indent = format!("\n{}", indent);
        let indent = indent.as_str();
        let body = statements.join(indent);
        format!("{}{}", declaration, body)
    } 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 = dump(py, indent, target);
        let iter = dump(py, indent, iter);

        let declaration = format!("for {} in {}:", target, iter);
        let mut statements = vec!(String::new());
        for statement in body.iter(py).unwrap() {
            let statement = dump(py, indent + 1, statement.unwrap());
            statements.push(statement);
        }
        let indent: String = iter::repeat("    ").take(indent + 1).collect();
        let indent = format!("\n{}", indent);
        let indent = indent.as_str();
        let body = statements.join(indent);
        format!("{}{}", declaration, body)
    } else if is_instance(&ast, &arguments_type) {
        let args_list = ast.getattr(py, "args").unwrap();
        let mut arguments = vec!();
        for arg in args_list.iter(py).unwrap() {
            let arg = arg.unwrap();
            arguments.push(dump(py, indent, arg));
        }
        format!("{}", arguments.join(", "))
    } else if is_instance(&ast, &arg_type) {
        let arg = ast.getattr(py, "arg").unwrap();
        format!("{}", arg)
    } 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 = dump(py, indent, left);
        let ops = ops.iter(py).unwrap();
        let comparators = comparators.iter(py).unwrap();

        let mut comparisons = vec!();
        for (op, comparator) in ops.zip(comparators) {
            let op = op.unwrap();
            let comparator = comparator.unwrap();

            let op = dump(py, indent, op);
            let comparator = dump(py, indent, comparator);

            comparisons.push(format!("{} {}", op, comparator));
        }
        format!("{} {}", left, comparisons.join(" "))
    } else if is_instance(&ast, &assign_type) {
        let targets = ast.getattr(py, "targets").unwrap();
        let value = ast.getattr(py, "value").unwrap();
        let mut cibles = vec!();
        for target in targets.iter(py).unwrap() {
            let target = target.unwrap();
            let target = dump(py, indent, target);
            cibles.push(target.to_string());
        }
        let value = dump(py, indent, value);
        format!("{} = {}", cibles.join(", "), value)
    } else if is_instance(&ast, &return_type) {
        let value = ast.getattr(py, "value").unwrap();
        let value = dump(py, indent, value);
        format!("return {}", value)
    } 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 = dump(py, indent, left);
        let op = dump(py, indent, op);
        let right = dump(py, indent, right);

        format!("{} {} {}", left, op, right)
    } else if is_instance(&ast, &name_type) {
        let id = ast.getattr(py, "id").unwrap();
        format!("{}", id)
    } else if is_instance(&ast, &name_constant_type) {
        let value = ast.getattr(py, "value").unwrap();
        format!("{}", value)
    } else if is_instance(&ast, &expr_type) {
        let value = ast.getattr(py, "value").unwrap();
        let value = dump(py, indent, value);
        format!("{}", value)
    } 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 = dump(py, indent, func);

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

        format!("{}({})", func, arguments.join(", "))
    } 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 mut arguments = vec!();
        for name in names.iter(py).unwrap() {
            let name = name.unwrap();
            arguments.push(dump(py, indent, name));
        }

        format!("from {} import {}", module, arguments.join(", "))
    } 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 = get_str(py, asname);

        if asname == "None" {
            format!("{}", name)
        } else {
            format!("{} as {}", name, asname)
        }
    } else if is_instance(&ast, &num_type) {
        let n = ast.getattr(py, "n").unwrap();
        format!("{}", n)
    } else if is_instance(&ast, &str_type) {
        let s = ast.getattr(py, "s").unwrap();
        format!("\"{}\"", s)
    } else if is_instance(&ast, &add_type) {
        format!("+")
    } else if is_instance(&ast, &mult_type) {
        format!("*")
    } else if is_instance(&ast, &eq_type) {
        format!("==")
    } else if is_instance(&ast, &lt_type) {
        format!("<")
    } else {
        format!("unknown {}", ast)
    }
}

#[allow(dead_code)]
pub fn dump_module(module: &PyObject) {
    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();
    for statement in body.iter(py).unwrap() {
        println!("{}", dump(py, 0, statement.unwrap()));
    }
}