generator: add intelligent whitespace based on tokens on nodes in a sequence

lib/6to5/generator.js

@@ -15,6 +15,7 @@ function CodeGenerator(code, ast, opts) {
   opts = opts || {}
 
   this.tabWidth = 2;
+  this.code     = code;
   this.opts     = opts;
   this.ast      = ast;
   this.buf      = "";
@@ -36,7 +37,7 @@ function CodeGenerator(code, ast, opts) {
 
 CodeGenerator.prototype.mark = function (node, locType) {
   var loc = node.loc;
-  if (!loc) return; // no locaiton info
+  if (!loc) return; // no location info
 
   var map = this.map;
   if (!map) return; // no source map
@@ -52,6 +53,7 @@ CodeGenerator.prototype.mark = function (node, locType) {
 };
 
 CodeGenerator.prototype.newline = function () {
+  this.buf = this.buf.replace(/\n(\s+)$/, "\n");
   this.buf += "\n";
   this.line++;
 };
@@ -126,28 +128,124 @@ CodeGenerator.prototype.generate = function () {
 CodeGenerator.prototype.buildPrint = function (parent) {
   var self = this;
 
-  var print = function (node) {
-    return self.print(node, parent);
+  var print = function (node, opts) {
+    return self.print(node, parent, opts);
   };
 
-  print.sequence = function (nodes) {
-    return self.printJoin(print, nodes);
+  print.sequence = function (nodes, opts) {
+    return self.printSequence(print, nodes, opts);
   };
 
   return print;
 };
 
-CodeGenerator.prototype.print = function (node, 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;
+
+    _.each(tokens, function (token, i) {
+      // this is the token this node starts with
+      if (node.start === token.start) {
+        startToken = tokens[i - 1];
+        endToken = token;
+        return false;
+      }
+    });
+
+    return hasWhitespaceBetween(startToken, endToken);
+  };
+
+  var needsNewlineAfter = function (node) {
+    var startToken;
+    var endToken;
+
+    _.each(tokens, function (token, i) {
+      // this is the token this node ends with
+      if (node.end === token.end) {
+        startToken = token;
+        endToken = tokens[i + 1];
+        return false;
+      }
+    });
+
+    return hasWhitespaceBetween(startToken, endToken);
+  };
+
+  opts.print = function (node, i) {
+    print(node, node && node.start != null && {
+      before: function () {
+        if (needsNewlineBefore(node)) self.newline();
+      },
+
+      after: function () {
+        if (needsNewlineAfter(node)) self.newline();
+      }
+    });
+  };
+
+  return this.printJoin(print, nodes, "\n", opts);
+};
+
+CodeGenerator.prototype.print = function (node, parent, opts) {
   if (!node) return "";
 
+  opts = opts || {};
+
   if (this[node.type]) {
     this.printLeadingComments(node);
     this.mark(node, "start");
 
+    if (opts.before) opts.before();
     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");
     this.printTrailingComments(node);
@@ -159,9 +257,7 @@ CodeGenerator.prototype.print = function (node, parent) {
 CodeGenerator.prototype.generateComment = function (comment) {
   var val = comment.value;
   if (comment.type === "Line") {
-    if (_.last(val) !== "\n") {
-      val += "\n";
-    }
+    if (_.last(val) !== "\n") val += "\n";
     return "//" + val;
   } else {
     return "/*" + val + "*/\n";
@@ -192,12 +288,16 @@ CodeGenerator.prototype.printJoin = function (print, nodes, sep, opts) {
   if (opts.indent) self.indent();
 
   _.each(nodes, function (node, i) {
-    if (opts.iterator) {
-      opts.iterator(node, i);
+    if (opts.print) {
+      opts.print(node, i);
     } else {
       print(node);
     }
 
+    if (opts.iterator) {
+      opts.iterator(node, i);
+    }
+
     if (i < len - 1) {
       self.push(sep);
     }