prettyprinter.java

Python Development Environment (Python IDE plugin for Eclipse). Features editor, code completion, re

/*
 * Created on Feb 11, 2006
 */
package org.python.pydev.parser.prettyprinter;


import java.io.IOException;
import java.util.Arrays;
import java.util.Collections;

import org.python.pydev.parser.jython.SimpleNode;
import org.python.pydev.parser.jython.SpecialStr;
import org.python.pydev.parser.jython.ast.Assert;
import org.python.pydev.parser.jython.ast.Assign;
import org.python.pydev.parser.jython.ast.Attribute;
import org.python.pydev.parser.jython.ast.AugAssign;
import org.python.pydev.parser.jython.ast.BinOp;
import org.python.pydev.parser.jython.ast.BoolOp;
import org.python.pydev.parser.jython.ast.Break;
import org.python.pydev.parser.jython.ast.Call;
import org.python.pydev.parser.jython.ast.ClassDef;
import org.python.pydev.parser.jython.ast.Compare;
import org.python.pydev.parser.jython.ast.Comprehension;
import org.python.pydev.parser.jython.ast.Continue;
import org.python.pydev.parser.jython.ast.Delete;
import org.python.pydev.parser.jython.ast.Dict;
import org.python.pydev.parser.jython.ast.Exec;
import org.python.pydev.parser.jython.ast.Expr;
import org.python.pydev.parser.jython.ast.For;
import org.python.pydev.parser.jython.ast.FunctionDef;
import org.python.pydev.parser.jython.ast.Global;
import org.python.pydev.parser.jython.ast.If;
import org.python.pydev.parser.jython.ast.IfExp;
import org.python.pydev.parser.jython.ast.Import;
import org.python.pydev.parser.jython.ast.ImportFrom;
import org.python.pydev.parser.jython.ast.Index;
import org.python.pydev.parser.jython.ast.Lambda;
import org.python.pydev.parser.jython.ast.ListComp;
import org.python.pydev.parser.jython.ast.Module;
import org.python.pydev.parser.jython.ast.Name;
import org.python.pydev.parser.jython.ast.NameTok;
import org.python.pydev.parser.jython.ast.Num;
import org.python.pydev.parser.jython.ast.Pass;
import org.python.pydev.parser.jython.ast.Print;
import org.python.pydev.parser.jython.ast.Raise;
import org.python.pydev.parser.jython.ast.Return;
import org.python.pydev.parser.jython.ast.Slice;
import org.python.pydev.parser.jython.ast.Str;
import org.python.pydev.parser.jython.ast.StrJoin;
import org.python.pydev.parser.jython.ast.Subscript;
import org.python.pydev.parser.jython.ast.TryExcept;
import org.python.pydev.parser.jython.ast.TryFinally;
import org.python.pydev.parser.jython.ast.Tuple;
import org.python.pydev.parser.jython.ast.UnaryOp;
import org.python.pydev.parser.jython.ast.While;
import org.python.pydev.parser.jython.ast.With;
import org.python.pydev.parser.jython.ast.Yield;
import org.python.pydev.parser.jython.ast.aliasType;
import org.python.pydev.parser.jython.ast.argumentsType;
import org.python.pydev.parser.jython.ast.commentType;
import org.python.pydev.parser.jython.ast.comprehensionType;
import org.python.pydev.parser.jython.ast.decoratorsType;
import org.python.pydev.parser.jython.ast.excepthandlerType;
import org.python.pydev.parser.jython.ast.exprType;
import org.python.pydev.parser.jython.ast.stmtType;
import org.python.pydev.parser.jython.ast.suiteType;
import org.python.pydev.parser.visitors.NodeUtils;

/**
 * statements that 'need' to be on a new line:
 * print
 * del
 * pass
 * flow
 * import
 * global
 * exec
 * assert
 * 
 * 
 * flow:
 * return
 * yield
 * raise
 * 
 * 
 * compound:
 * if
 * while
 * for
 * try
 * func
 * class
 * 
 * @author Fabio
 */
public class PrettyPrinter extends PrettyPrinterUtils{

    /**
     * If this is true, we don't add a new-line after the statement (when we would normally
     * add a new line for the next statement).
     */
    private boolean isSingleStmt;

    @Override
    protected boolean fixNewStatementCondition() throws IOException {
        boolean ret = false;
        if(!isSingleStmt){
            ret = super.fixNewStatementCondition();
        }
        return ret;
    }

    public PrettyPrinter(PrettyPrinterPrefs prefs, IWriterEraser writer){
        this(prefs, writer, false);
    }
    
    public PrettyPrinter(PrettyPrinterPrefs prefs, IWriterEraser writer, boolean isSingleStmt){
        this.prefs = prefs;
        this.isSingleStmt = isSingleStmt;
        state = new WriteState(writer, prefs);
        auxComment = new AuxSpecials(state, prefs);
    }
    

    @Override
    public Object visitModule(Module node) throws Exception {
        super.visitModule(node);
        if(node.specialsAfter != null){
            for(Object o :node.specialsAfter){
                commentType t = (commentType) o;
                String c = t.id.trim();
                state.write(c);
            }
        }
        return null;
    }
    
    @Override
    public Object visitImportFrom(ImportFrom node) throws Exception {
        auxComment.writeSpecialsBefore(node);
        
        
        auxComment.writeSpecialsBefore(node.module, new String[0], new String[0], true);
        auxComment.writeSpecialsBefore(node.module, null, null, false);
        
        state.write(((NameTok)node.module).id);
        auxComment.writeSpecialsAfter(node.module);
        
        for (aliasType name : node.names){
            auxComment.writeSpecialsBefore(name);
            name.accept(this);
            auxComment.writeSpecialsAfter(name);
        }
        afterNode(node);
        fixNewStatementCondition();
        return null;
    }
    
    @Override
    public Object visitImport(Import node) throws Exception {
        auxComment.writeSpecialsBefore(node);
        
        for (aliasType name : node.names){
            auxComment.writeSpecialsBefore(name);
            name.accept(this);
            auxComment.writeSpecialsAfter(name);
        }
        afterNode(node);
        return null;
    }

    @Override
    public Object visitAssign(Assign node) throws Exception {
    	state.pushInStmt(node);
        auxComment.writeSpecialsBefore(node);
        for (int i = 0; i < node.targets.length; i++) {
            exprType target = node.targets[i];
            if(i == node.targets.length -1){
                //last one: cannot have comments, if it has, we will have to move them to the next node.
                auxComment.moveComments(target, node.value, false, true);
            }
            if(i >= 1){ //more than one assign
                state.write(" = ");
            }
            target.accept(this);
        }
        state.write(" = ");
        
        node.value.accept(this);

        checkEndRecord();
        if(auxComment.hasCommentsAfter(node)){
            auxComment.writeSpecialsAfter(node, false);
            checkEndRecord();
        }
        
        state.popInStmt();
        fixNewStatementCondition();
        return null;
    }
    
    public void visitBeforeLeft(SimpleNode node) throws IOException{
        auxComment.writeSpecialsBefore(node);
        state.pushInStmt(node);
    }
    public void visitAfterRight(SimpleNode node) throws IOException{
        state.popInStmt();
        
        auxComment.writeSpecialsAfter(node, false);
        
        if(!state.inStmt()){
            checkEndRecord();
        }
        
    }
    
    @Override
    public Object visitAugAssign(AugAssign node) throws Exception {
        visitBeforeLeft(node);
        node.target.accept(this);
        state.write(augOperatorMapping[node.op]);
        node.value.accept(this);
        visitAfterRight(node);
        return null;
    }
    
    @Override
    public Object visitBinOp(BinOp node) throws Exception {
        visitBeforeLeft(node);
        node.left.accept(this);
        state.write(operatorMapping[node.op]);
        node.right.accept(this);
        visitAfterRight(node);
        return null;
    }
    

    @Override
    public Object visitUnaryOp(UnaryOp node) throws Exception {
        visitBeforeLeft(node);
        state.write(unaryopOperatorMapping[node.op]);
        node.operand.accept(this);
        visitAfterRight(node);
        return null;
    }
    

    @Override
    public Object visitBoolOp(BoolOp node) throws Exception {
        auxComment.writeSpecialsBefore(node);
        state.pushInStmt(node);
        for (int i = 0; i < node.values.length-1; i++) {
            node.values[i].accept(this);
            state.write(boolOperatorMapping[node.op]);
        }
        node.values[node.values.length-1].accept(this);
        visitAfterRight(node);
        return null;
    }
    
    @Override
    public Object visitSubscript(Subscript node) throws Exception {
        node.value.accept(this);
        
        visitBeforeLeft(node);
        node.slice.accept(this);
        visitAfterRight(node);
        return null;
    }

    @Override
    public Object visitCompare(Compare node) throws Exception {
        auxComment.writeSpecialsBefore(node);
        node.left.accept(this);
        
        for (int i = 0; i < node.comparators.length; i++) {
            state.write(cmpop[node.ops[i]]);
            node.comparators[i].accept(this);
        }
        
        auxComment.writeSpecialsAfter(node);
        return null;
    }
    
    
    @Override
    public Object visitDict(Dict node) throws Exception {
        state.indent();
        auxComment.writeSpecialsBefore(node);
        exprType[] keys = node.keys;
        exprType[] values = node.values;
        for (int i = 0; i < values.length; i++) {
            prefs.enableSpacesAfterColon();
            keys[i].accept(this);
            values[i].accept(this);
            prefs.disableSpacesAfterColon();
        }
        auxComment.writeSpecialsAfter(node);
        dedent();
        return null;
    }
    
    @Override
    public Object visitLambda(Lambda node) throws Exception {
        genericBefore(node, false);
        state.pushInStmt(node);
        prefs.enableSpacesAfterColon();
        makeArgs(node.args.args, node.args);
        prefs.disableSpacesAfterColon();
        node.body.accept(this);
        state.popInStmt();
        genericAfter(node, false, false);
        return null;
    }
    
    @Override
    public Object visitList(org.python.pydev.parser.jython.ast.List node) throws Exception{
        return visitGeneric(node, "visitList", false, null, true);
    }

    @Override
    public Object visitDelete(Delete node) throws Exception {
    	return visitGeneric(node, "visitDelete", false);
    }

    @Override
    public Object visitListComp(ListComp node) throws Exception {
        beforeNode(node);
        node.elt.accept(this);
        for(comprehensionType c:node.generators){
            c.accept(this);
        }
        afterNode(node);
    	return null;
    }

    private SimpleNode[] reverseNodeArray(SimpleNode[] expressions) {
        java.util.List<SimpleNode> ifs = Arrays.asList(expressions);
        Collections.reverse(ifs);
        SimpleNode[] ifsInOrder = ifs.toArray(new SimpleNode[0]);
        return ifsInOrder;
    }

    @Override
    public Object visitComprehension(Comprehension node) throws Exception {
        beforeNode(node);
        node.target.accept(this);
        node.iter.accept(this);
        for(SimpleNode s:reverseNodeArray(node.ifs)){
            s.accept(this);
        }
        afterNode(node);
        return null;
    }
    
    @Override
    public Object visitExpr(Expr node) throws Exception {
        return visitGeneric(node, "visitExpr", false, null, false, false);
    }
    
    @Override
    public Object visitWhile(While node) throws Exception {
        auxComment.writeSpecialsBefore(node);
        state.indent();
        state.pushInStmt(node.test);
        node.test.accept(this);
        state.popInStmt();
        afterNode(node);
        fixNewStatementCondition();
        for(SimpleNode n: node.body){
            n.accept(this);
        }
        dedent();
        
        if(node.orelse != null){
            state.indent();
            auxComment.writeSpecialsBefore(node.orelse);
            afterNode(node.orelse);
            node.orelse.accept(this);
            dedent();
        }
        return null;
    }
    
    
    @Override
    public Object visitFor(For node) throws Exception {
        //for a in b: xxx else: yyy
        
        state.pushInStmt(node);
        //for
        auxComment.writeSpecialsBefore(node);
        state.indent();
        
        //a
        node.target.accept(this);
        
        //in b
        state.pushInStmt(node.iter);
        node.iter.accept(this);
        state.popInStmt();
        afterNode(node);
        state.popInStmt();
        
        fixNewStatementCondition();
        for(SimpleNode n: node.body){
            n.accept(this);
        }
        dedent();
        
        if(node.orelse != null){
            state.indent();
            auxComment.writeSpecialsBefore(node.orelse);
            auxComment.writeSpecialsAfter(node.orelse);
            afterNode(node.orelse);
            node.orelse.accept(this);
            dedent();
        }
        return null;
    }


    @Override
    public Object visitTuple(Tuple node) throws Exception {
        visitGeneric(node, "visitTuple", false, null, true);
        return null;
    }

    
    
    
    @Override
    public Object visitRaise(Raise node) throws Exception {
        visitGeneric(node, "visitRaise", true);