rewritervisitor.java

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

package org.python.pydev.refactoring.ast.rewriter;

import java.io.StringWriter;
import java.io.Writer;
import java.util.Arrays;
import java.util.Collections;
import java.util.Iterator;

import org.python.pydev.parser.jython.SimpleNode;
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.Ellipsis;
import org.python.pydev.parser.jython.ast.Exec;
import org.python.pydev.parser.jython.ast.Expr;
import org.python.pydev.parser.jython.ast.Expression;
import org.python.pydev.parser.jython.ast.ExtSlice;
import org.python.pydev.parser.jython.ast.For;
import org.python.pydev.parser.jython.ast.FunctionDef;
import org.python.pydev.parser.jython.ast.GeneratorExp;
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.Interactive;
import org.python.pydev.parser.jython.ast.Lambda;
import org.python.pydev.parser.jython.ast.List;
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.Repr;
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.Suite;
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.decoratorsType;
import org.python.pydev.parser.jython.ast.excepthandlerType;
import org.python.pydev.parser.jython.ast.exprType;
import org.python.pydev.parser.jython.ast.keywordType;
import org.python.pydev.parser.jython.ast.listcompType;
import org.python.pydev.parser.jython.ast.modType;
import org.python.pydev.parser.jython.ast.stmtType;
import org.python.pydev.parser.jython.ast.suiteType;
import org.python.pydev.parser.prettyprinter.PrettyPrinter;
import org.python.pydev.parser.prettyprinter.PrettyPrinterPrefs;
import org.python.pydev.parser.prettyprinter.WriterEraser;
import org.python.pydev.refactoring.ast.printer.SourcePrinter;
import org.python.pydev.refactoring.ast.visitors.VisitorFactory;

/**
 * GoF Visitor Pattern for AST traversal (see AbstractRewriterVisitor for implemented Jython Visitor class). It will traverse an Abstract
 * Syntax Tree, in order to reprint the source code.
 * 
 * @author Ueli Kistler (ukistler[at]hsr.ch)
 * 
 */
public class RewriterVisitor extends AbstractRewriterVisitor {
    
    //------------------------------------------------------------------------------------------------- public interface

    //REWRITER INTERFACE FOR CURRENT VERSION ---------------------------------------------------------------------------
    public static String reparsed(String source, String string) {
        StringWriter out = new StringWriter();
        createRewriterVisitor(out, source, "\n");
        return out.getBuffer().toString();
    }

    public static String createSourceFromAST(SimpleNode root, boolean ignoreComments, String newLineDelim) {
        RewriterVisitor visitor = null;
        StringWriter writer = new StringWriter();
        try {
            visitor = new RewriterVisitor(VisitorFactory.createPrinter(writer, newLineDelim));
            visitor.setIgnoreComments(ignoreComments);
            visitor.visit(root);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
        visitor.flush();
        return writer.getBuffer().toString();
    }

    public static String createSourceFromAST(SimpleNode root, String newLineDelim) {
        return createSourceFromAST(root, false, newLineDelim);
    }

    private static RewriterVisitor createRewriterVisitor(Writer out, String source, String newLineDelim) {
        RewriterVisitor visitor = null;
        try {
            SimpleNode root = VisitorFactory.getRootNodeFromString(source);
            visitor = new RewriterVisitor(VisitorFactory.createPrinter(out, newLineDelim));
            root.accept(visitor);
        } catch (Throwable e) {
            throw new RuntimeException(e);
        }
        visitor.flush();
        return visitor;
    }
    
    //END REWRITER INTERFACE FOR CURRENT VERSION -----------------------------------------------------------------------


//    public static String reparsed(String source, String delimiter) {
//        try {
//            SimpleNode root = VisitorFactory.getRootNodeFromString(source);
//            return createSourceFromAST(root, delimiter);
//        } catch (Throwable e) {
//            throw new RuntimeException(e);
//        }
//        
//    }
//    
//    public static String createSourceFromAST(SimpleNode root, String newLineDelim) {
//        return createSourceFromAST(root, false, newLineDelim);
//    }
//
//
//    public static String createSourceFromAST(SimpleNode root, boolean ignoreComments, String newLineDelim) {
//        final WriterEraser stringWriter = new WriterEraser();
//        PrettyPrinterPrefs prettyPrinterPrefs = new PrettyPrinterPrefs(newLineDelim);
//        prettyPrinterPrefs.setSpacesAfterComma(1);
//        prettyPrinterPrefs.setSpacesBeforeComment(1);
//        PrettyPrinter printer = new PrettyPrinter(prettyPrinterPrefs, stringWriter);
//        try {
//            root.accept(printer);
//            return stringWriter.getBuffer().toString();
//        } catch (Exception e) {
//            throw new RuntimeException(e);
//        }
//        
//    }


    
    //--------------------------------------------------------------------------------------------- end public interface


	private RewriterVisitor(SourcePrinter printer) {
		super(printer);
	}

	private SimpleNode handleCallArguments(Call node) throws Exception {
		SimpleNode lastNode = null;
		lastNode = visitWithSeparator(node, node.args);
		if (printer.getNodeHelper().isFilledList(node.keywords)) {
			// comma before already handeld in visitWithSeparator
			lastNode = visitWithSeparator(node, node.keywords);
		}
		if (node.starargs != null) {
			// comma before already handeld in visitWithSeparator
			lastNode = visit(node.starargs);
		}
		if (node.kwargs != null) {
			if (node.starargs != null)
				printer.printListSeparator();
			lastNode = visit(node.kwargs);
		}
		return lastNode;
	}

	private boolean handleCommaOptional(SimpleNode node) {
		return (printer.hasSpecialAfter(node, printer.getSyntaxhelper().getComma()));
	}

	protected void handleCommentAfter(SimpleNode node) {
		printer.printCommentAfter(node);
	}

	private void handleCommentAfterBody(SimpleNode node, SimpleNode firstBodyNode) {
		printer.printCommentAfterBody(node, firstBodyNode);
	}

	protected void handleCommentBefore(SimpleNode node) {
		printer.printCommentBefore(node);
	}

	private void handleCommentBeforeBody(SimpleNode node, SimpleNode firstBodyNode) {
		printer.printCommentBeforeBody(node, firstBodyNode);
	}

	private SimpleNode handleDecoratorArgs(decoratorsType node) throws Exception {
		if (node == null)
			return null;
		SimpleNode lastNode = null;

		if (node.args != null) {
			if (isFilledList(node.args))
				printer.printBeforeTuple();

			lastNode = visitWithSeparator(node, reverseNodeArray(node.args));

			if (isFilledList(node.keywords)) {
				printer.printListSeparator();
				lastNode = visitWithSeparator(node, reverseNodeArray(node.keywords));
			}
			if (node.starargs != null) {
				// comma before already handeld in visitWithSeparator
				printer.printBeforeVarArg();
				super.visit(node.starargs);
				lastNode = node.starargs;
			}
			if (node.kwargs != null) {
				if (lastNode != null)
					printer.printListSeparator();
				printer.printBeforeKwArg();
				super.visit(node.kwargs);
				lastNode = node.kwargs;
			}

			if (isFilledList(node.args))
				printer.printAfterTuple();

		}
		return lastNode;
	}

	private SimpleNode handleFunctionArgs(SimpleNode node, argumentsType arguments) throws Exception {
		SimpleNode lastNode = null;

		if (arguments != null) {
			lastNode = handlePositionalAndDefaultArgs(node, arguments);
			lastNode = handleVarArgs(arguments, lastNode);
			lastNode = handleKwArgs(arguments, lastNode);
		}
		return lastNode;
	}

	private void handleIfElse(If node) throws Exception {
		boolean passedElsif = false;
		if (node.orelse != null) {
			for (int i = 0; i < node.orelse.length; i++) {
				passedElsif = handleIfElseSuite(node, passedElsif, node.orelse[i]);

			}
			if (passedElsif)
				printer.outdent();

		}
	}

	private boolean handleIfElseSuite(If node, boolean passedElsif, stmtType statement) throws Exception {
		printer.setDisabledIfPrinting(true);

		printer.printNewlineAndIndentation();

		if (printer.getNodeHelper().isIfStatement(statement)) {
			if (!passedElsif) {
				printer.printStatementElif();
			} else
				printer.setDisabledIfPrinting(false);
		} else if (!passedElsif) {
			passedElsif = true;
			printer.printStatementElse();
			printer.printFunctionMarker();
			printer.indent();
			printer.printNewlineAndIndentation();
		}

		visit(statement);
		return passedElsif;
	}

	private exprType handleKeyValueArgs(SimpleNode parent, exprType[] keys, exprType[] values) throws Exception {
		exprType lastNode = null;
		if (keys == null)
			return lastNode;

		int startOffset = 0;

		if (values != null)
			startOffset = (keys.length - values.length);

		for (int i = 0; i < keys.length; i++) {
			super.visit(keys[i]);
			if ((values != null) && (i >= startOffset)) {
				if (values[i - startOffset] != null) {
					if (printer.getNodeHelper().isDict(parent)) {
						printer.printDictBeforeValue();
					} else {
						printer.printAssignmentOperator(false, false);
					}
					super.visit(values[i - startOffset]);
				}

			}
			if (i < keys.length - 1)
				printer.printListSeparator();
			else
				lastNode = keys[i];

		}
		return lastNode;
	}

	private SimpleNode handleKwArgs(argumentsType arguments, SimpleNode lastNode) throws Exception {
		if (arguments.kwarg != null) {
			if (lastNode != null)
				printer.printListSeparator();
			printer.printBeforeKwArg();
			super.visit(arguments.kwarg);
			lastNode = arguments.kwarg;

			if (handleCommaOptional(lastNode))
				printer.printListSeparator();
		}
		return lastNode;
	}

	private SimpleNode handlePositionalAndDefaultArgs(SimpleNode node, argumentsType arguments) throws Exception {
		SimpleNode lastNode;
		lastNode = visitKeyValue(node, arguments.args, arguments.defaults);
		if (handleCommaOptional(lastNode) && arguments.vararg == null && arguments.kwarg == null)
			printer.printListSeparator();
		return lastNode;
	}

	private void handlePostNode(SimpleNode parent, SimpleNode lastNode, Iterator<SimpleNode> iter, boolean outdent, boolean separator) {
		if (separator)
			handleSeparator(parent, lastNode, iter);
		handleCommentAfter(lastNode);

		if (outdent) {
			if (!iter.hasNext()) {
				printer.outdent();
			}
		}
		if (iter.hasNext() && !separator) {
			if (!(printer.getNodeHelper().isComprehension(parent))) {
				printer.printNewlineAndIndentation();
			}
		}
	}

	private void handlePreNode(SimpleNode parent, SimpleNode lastNode, java.util.List<SimpleNode> nodes) {
		handleCommentBefore(lastNode);
		if (printer.getNodeHelper().isComprehension(parent)) {
			printer.printStatementIf(lastNode, true, true);
		}
	}

	private void handleRootNode(modType node, stmtType[] body) throws Exception {
		handleCommentBefore(node);
		handleCommentAfter(visit(node, body, false, false));
		handleCommentAfter(node);
	}

	private void handleSeparator(SimpleNode parent, SimpleNode lastNode, Iterator<SimpleNode> iter) {
		if (iter.hasNext()) {
			if (printer.getNodeHelper().isBoolOp(parent)) {
				BoolOp boolParent = (BoolOp) parent;
				printer.printBoolOp(boolParent.op);
			} else {
				printer.printListSeparator();
			}
		} else {
			if (handleCommaOptional(lastNode)) {
				printer.printListSeparator();
			}

		}