view src/ast_dump.rs @ 14:719a27f1c1c7

Add ast.ListComp
author Bastien Orivel <eijebong@bananium.fr>
date Thu, 02 Jun 2016 20:40:24 +0200
parents 38b0d63697b1
children b21a246349f3
line wrap: on
line source

use python_ast::{Module, stmt, expr, boolop, operator, unaryop, cmpop, arguments, arg, keyword, comprehension};

use std::iter;

impl boolop {
    fn to_string(&self) -> &'static str {
        match *self {
            boolop::And => "and",
            boolop::Or => "or"
        }
    }
}

impl operator {
    fn to_string(&self) -> &'static str {
        match *self {
            operator::Add => "+",
            operator::Sub => "-",
            operator::Mult => "*",
            operator::MatMult => "@",
            operator::Div => "/",
            operator::Mod => "%",
            operator::Pow => "**",
            operator::LShift => "<<",
            operator::RShift => ">>",
            operator::BitOr => "|",
            operator::BitXor => "^",
            operator::BitAnd => "&",
            operator::FloorDiv => "//",
        }
    }
}

impl unaryop {
    fn to_string(&self) -> &'static str {
        match *self {
            unaryop::Invert => "~",
            unaryop::Not => "!",
            unaryop::UAdd => "+",
            unaryop::USub => "-"
        }
    }
}

impl cmpop {
    fn to_string(&self) -> &'static str {
        match *self {
            cmpop::Eq => "==",
            cmpop::NotEq => "!=",
            cmpop::Lt => "<",
            cmpop::LtE => "<=",
            cmpop::Gt => ">",
            cmpop::GtE => ">=",
            cmpop::Is => "is",
            cmpop::IsNot => "is not",
            cmpop::In => "in",
            cmpop::NotIn => "not in",
        }
    }
}

fn args_to_string(args: Vec<expr>) -> String {
    let mut arguments = vec!();
    for arg in args {
        arguments.push(arg.to_string());
    }
    arguments.join(", ")
}

fn kwargs_to_string(kwargs: Vec<keyword>) -> String {
    let mut arguments = vec!();
    for arg in kwargs {
        match arg.arg {
            Some(arg) => arguments.push(arg),
            None => arguments.push("**kwarg".to_string())
        }
    }
    arguments.join(", ")
}

fn arguments_to_string(args: arguments) -> String {
    let mut arguments = vec!();
    for arg in args.args {
        arguments.push(arg.arg);
    }
    arguments.join(", ")
}

fn statements_to_string(indent: usize, body: Vec<stmt>) -> String {
    let mut statements = vec!();
    for statement in body {
        statements.push(statement.to_string(indent + 1));
    }
    statements.join("\n")
}

fn if_else_statements_to_string(indent: usize, body: Vec<stmt>, orelse: Vec<stmt>) -> String {
    let body = statements_to_string(indent, body);
    if orelse.is_empty() {
        body
    } else {
        format!("{}\n{}else:\n{}", body, make_indent(indent), statements_to_string(indent, orelse))
    }
}

fn generators_to_string(generators: Vec<comprehension>) -> String {
    let mut result  = vec!();
    for generator in generators {
        result.push(format!("for {} in {}", generator.target.to_string(), generator.iter.to_string()));
        for if_ in generator.ifs {
            result.push(format!("if {}", if_.to_string()))
        }
    }
    result.join(" ")
}
impl expr {
    fn to_string(&self) -> String {
        match self.clone() {
            expr::BoolOp(op, values) => format!("{}({})", op.to_string(), args_to_string(values)),
            expr::BinOp(a, op, b) => format!("{} {} {}", a.to_string(), op.to_string(), b.to_string()),
            expr::UnaryOp(op, operand) => format!("{}{}", op.to_string(), operand.to_string()),

            expr::Compare(left, ops, comparators) => format!("{} {}", left.to_string(), {
                let mut arguments = vec!();

                // XXX: wrong order!
                for op in ops {
                    arguments.push(op.to_string().to_string())
                }
                for comparator in comparators {
                    arguments.push(comparator.to_string())
                }
                /*
                for (op, comparator) in ops.zip(comparators) {
                    let op = op.unwrap();
                    let comparator = comparator.unwrap();
                    arguments.push(format!("{} {}", op.to_string(), comparator.to_string()))
                }
                */

                arguments.join(" ")
            }),
            expr::Call(func, args, keywords) => format!("{}({}{})", func.to_string(), args_to_string(args), kwargs_to_string(keywords)),
            expr::Num(n) => format!("{}", n),
            expr::Str(s) => format!("\"{}\"", s),
            expr::NameConstant(name) => format!("{}", name),
            expr::Attribute(value, attr, ctx) => format!("{}.{}", value.to_string(), attr),
            expr::Name(name, ctx) => format!("{}", name),
            expr::List(elts, ctx) => format!("[{}]", args_to_string(elts)),
            expr::ListComp(elt, generators) => format!("[{} {}]", elt.to_string(), generators_to_string(generators))
        }
    }
}

impl arguments {
    fn to_string(&self) -> String {
        let mut args = vec!();
        for arg in self.args.clone() {
            args.push(arg.arg);
        }
        args.join(", ")
    }
}

fn make_indent(indent: usize) -> String {
    iter::repeat("    ").take(indent).collect()
}

impl stmt {
    fn to_string(&self, indent: usize) -> String {
        let current_indent = make_indent(indent);
        match self.clone() {
            stmt::ClassDef(name, bases, keywords, body, decorator_list) => format!("{}def {}({}):\n{}", current_indent, name, args_to_string(bases), statements_to_string(indent, body)),
            stmt::FunctionDef(name, arguments, body, decorator_list, returns) => format!("{}def {}({}):\n{}", current_indent, name, arguments_to_string(arguments), statements_to_string(indent, body)),
            stmt::Global(names) => format!("{}global {}", current_indent, names.join(", ")),
            stmt::If(test, body, orelse) => format!("{}if {}:\n{}", current_indent, test.to_string(), if_else_statements_to_string(indent, body, orelse)),
            stmt::While(test, body, orelse) => format!("{}while {}:\n{}", current_indent, test.to_string(), if_else_statements_to_string(indent, body, orelse)),
            stmt::For(target, iter, body, orelse) => format!("{}for {} in {}:\n{}", current_indent, target.to_string(), iter.to_string(), if_else_statements_to_string(indent, body, orelse)),
            stmt::Assign(targets, value) => format!("{}{} = {}", current_indent, args_to_string(targets), value.to_string()),
            stmt::AugAssign(target, op, value) => format!("{}{} {}= {}", current_indent, target.to_string(), op.to_string(), value.to_string()),
            stmt::Return(expr) => format!("{}return{}", current_indent, match expr {
                Some(expr) => format!(" {}", expr.to_string()),
                None => "".to_string()
            }),
            stmt::ImportFrom(module, names, level) => format!("{}from {} import {}", current_indent, module.to_string(), {
                /*
                let mut exprs = vec!();
                for alias in names.iter() {
                    let alias = alias.to_string();
                    exprs.push(alias);
                }
                exprs.join(", ")
                */
                "".to_string()
            }),
            stmt::Expr(expr) => format!("{}{}", current_indent, expr.to_string()),
            stmt::Break => format!("{}break", current_indent)
        }
    }
}

#[allow(dead_code)]
pub fn dump_ast(ast: &Module) {
    for statement in &ast.statements {
        println!("{}", statement.to_string(0));
    }
}