generator: add _push method to simplify buffer pushing and location tracking

lib/6to5/generator.js

@@ -54,13 +54,11 @@ CodeGenerator.prototype.mark = function (node, locType) {
 
 CodeGenerator.prototype.newline = function () {
   this.buf = this.buf.replace(/\n(\s+)$/, "\n");
-  this.buf += "\n";
+  this._push("\n");
-  this.line++;
 };
 
 CodeGenerator.prototype.semicolon = function () {
-  this.buf += ";";
+  this._push(";");
-  this.column++;
 };
 
 CodeGenerator.prototype.keyword = function (name) {
@@ -69,7 +67,7 @@ CodeGenerator.prototype.keyword = function (name) {
 };
 
 CodeGenerator.prototype.push = function (str, noIndent) {
-  if (this._indent && !noIndent) {
+  if (this._indent && !noIndent && str !== "\n") {
     // we have an indent level and we aren't pushing a newline
     var indent = this.getIndent();
 
@@ -80,6 +78,10 @@ CodeGenerator.prototype.push = function (str, noIndent) {
     if (this.isLast("\n")) str = indent + str;
   }
 
+  this._push(str);
+};
+
+CodeGenerator.prototype._push = function (str) {
   var self = this;
 
   _.each(str, function (cha) {
@@ -140,50 +142,11 @@ CodeGenerator.prototype.buildPrint = function (parent) {
 };
 
 CodeGenerator.prototype.printSequence = function (print, nodes, opts) {
-  var comments = this.ast.comments;
   var tokens   = this.ast.tokens;
 
   var self = this;
   opts = opts || {};
 
-  // calculate the whitespace between two tokens
-  var hasWhitespaceBetween = function (startToken, endToken) {
-    if (!startToken) return false;
-
-    var whitespace = false;
-
-    var start = startToken.end;
-    var end   = endToken.start;
-
-    var sep = self.code.slice(start, end);
-
-    var lines = 0;
-    lines--; // take off the current line
-
-    // remove comments
-    _.each(comments, function (comment) {
-      // this comment is after the last token or befor ethe first
-      if (comment.end > end || comment.start < start) return;
-
-      var length = comment.end - comment.start;
-
-      // compute the relative positions of the comment within the sliced node
-      // string
-      var startRelative = comment.start - start;
-      var endRelative = comment.end - start;
-
-      // remove the line this comment ends with
-      if (comment.type === "Line") lines--;
-
-      // remove comment
-      sep = sep.slice(0, startRelative) + util.repeat(length) + sep.slice(endRelative);
-    });
-
-    // check if there was a newline between the two nodes
-    lines += _.size(sep.match(/\n/g));
-    return lines > 0;
-  };
-
   var needsNewlineBefore = function (node) {
     var startToken;
     var endToken;
@@ -197,7 +160,7 @@ CodeGenerator.prototype.printSequence = function (print, nodes, opts) {
       }
     });
 
-    return hasWhitespaceBetween(startToken, endToken);
+    return self.hasWhitespaceBetween(startToken, endToken);
   };
 
   var needsNewlineAfter = function (node) {
@@ -213,7 +176,7 @@ CodeGenerator.prototype.printSequence = function (print, nodes, opts) {
       }
     });
 
-    return hasWhitespaceBetween(startToken, endToken);
+    return self.hasWhitespaceBetween(startToken, endToken);
   };
 
   opts.print = function (node, i) {
@@ -231,6 +194,43 @@ CodeGenerator.prototype.printSequence = function (print, nodes, opts) {
   return this.printJoin(print, nodes, "\n", opts);
 };
 
+CodeGenerator.prototype.hasWhitespaceBetween = function (startToken, endToken) {
+  if (!endToken || !this.opts.whitespace) return false;
+
+  var comments = this.ast.comments;
+
+  var start = startToken ? startToken.end : 0;
+  var end   = endToken.start;
+
+  var sep = this.code.slice(start, end);
+
+  var lines = 0;
+  if (start > 0) lines--; // take off the current line
+
+  // remove comments
+  _.each(comments, function (comment) {
+    // this comment is after the last token or befor ethe first
+    if (comment.end > end || comment.start < start) return;
+
+    var length = comment.end - comment.start;
+
+    // compute the relative positions of the comment within the sliced node
+    // string
+    var startRelative = comment.start - start;
+    var endRelative = comment.end - start;
+
+    // remove the line this comment ends with
+    if (comment.type === "Line") lines--;
+
+    // remove comment
+    sep = sep.slice(0, startRelative) + util.repeat(length) + sep.slice(endRelative);
+  });
+
+  // check if there was a newline between the two nodes
+  lines += _.size(sep.match(/\n/g));
+  return lines > 0;
+};
+
 CodeGenerator.prototype.print = function (node, parent, opts) {
   if (!node) return "";
 
@@ -238,16 +238,19 @@ CodeGenerator.prototype.print = function (node, parent, opts) {
 
   if (this[node.type]) {
     this.printLeadingComments(node);
-    this.mark(node, "start");
 
     if (opts.before) opts.before();
+    this.mark(node, "start");
+
     var needsParans = t.needsParans(node, parent);
     if (needsParans) this.push("(");
+
     this[node.type](node, this.buildPrint(node), parent);
     if (needsParans) this.push(")");
-    if (opts.after) opts.after();
 
     this.mark(node, "end");
+    if (opts.after) opts.after();
+
     this.printTrailingComments(node);
   } else {
     throw new ReferenceError("unknown node " + node.type + " " + JSON.stringify(node));
@@ -257,24 +260,38 @@ CodeGenerator.prototype.print = function (node, parent, opts) {
 CodeGenerator.prototype.generateComment = function (comment) {
   var val = comment.value;
   if (comment.type === "Line") {
-    if (_.last(val) !== "\n") val += "\n";
+    val = "//" + val + "\n";
-    return "//" + val;
   } else {
-    return "/*" + val + "*/\n";
+    val = "/*" + val + "*/\n";
   }
+  return val;
 };
 
 CodeGenerator.prototype.printTrailingComments = function (node) {
-  var self = this;
+  this._printComments(node.trailingComments);
-  _.each(node.trailingComments, function (comment) {
-    self.push(self.generateComment(comment));
-  });
 };
 
 CodeGenerator.prototype.printLeadingComments = function (node) {
+  this._printComments(node.leadingComments);
+};
+
+CodeGenerator.prototype._printComments = function (comments) {
+  if (!comments || !comments.length) return;
+
   var self = this;
-  _.each(node.leadingComments, function (comment) {
+
+  _.each(comments, function (comment, i) {
+    // whitespace before
+    if (self.hasWhitespaceBetween(comment, comments[i - 1])) {
+      self.newline();
+    }
+
     self.push(self.generateComment(comment));
+
+    // whitespace after
+    if (self.hasWhitespaceBetween(comment, comments[i + 1])) {
+      self.newline();
+    }
   });
 };