Add support for preserving comments in babel-template. (#3689)

* Add support for preserving comments in babel-template.

* Add an API section to babel-template.

packages/babel-template/README.md

@@ -32,3 +32,33 @@ console.log(generate(ast).code);
 ```js
 var myModule = require('my-module');
 ```
+
+## API
+
+### `template(code, [opts])`
+
+#### code
+
+Type: `string`
+
+#### options
+
+`babel-template` accepts all of the options from [babylon], and specifies
+some defaults of its own:
+
+* `allowReturnOutsideFunction` is set to `true` by default.
+* `allowSuperOutsideMethod` is set to `true` by default.
+
+##### preserveComments
+
+Type: `boolean`
+Default: `false`
+
+Set this to `true` to preserve any comments from the `code` parameter.
+
+#### Return value
+
+`babel-template` returns a `function` which is invoked with an optional object
+of replacements. See the usage section for an example.
+
+[babylon]: https://github.com/babel/babylon#options

packages/babel-template/src/index.js

@@ -24,16 +24,19 @@ export default function (code: string, opts?: Object): Function {
     }
   }
 
+  opts = assign({
+    allowReturnOutsideFunction: true,
+    allowSuperOutsideMethod: true,
+    preserveComments: false,
+  }, opts);
+
   let getAst = function () {
     let ast;
 
     try {
-      ast = babylon.parse(code, assign({
-        allowReturnOutsideFunction: true,
-        allowSuperOutsideMethod: true
-      }, opts));
+      ast = babylon.parse(code, opts);
 
-      ast = traverse.removeProperties(ast);
+      ast = traverse.removeProperties(ast, {preserveComments: opts.preserveComments});
 
       traverse.cheap(ast, function (node) {
         node[FROM_TEMPLATE] = true;

packages/babel-template/test/index.js

@@ -1,6 +1,9 @@
+var generator = require('../../babel-generator').default;
 var template  = require("../lib");
 var chai      = require("chai");
 
+var comments = "// Sum two numbers\nconst add = (a, b) => a + b;";
+
 suite("templating", function () {
   test("import statement will cause parser to throw by default", function () {
     chai.expect(function () {
@@ -13,4 +16,15 @@ suite("templating", function () {
       template("import foo from 'foo'", {sourceType: 'module'})({});
     }).not.to.throw();
   });
+
+  test("should strip comments by default", function () {
+    var code = "const add = (a, b) => a + b;"
+    var output = template(comments)();
+    chai.expect(generator(output).code).to.be.equal(code);
+  });
+
+  test("should preserve comments with a flag", function () {
+    var output = template(comments, {preserveComments: true})();
+    chai.expect(generator(output).code).to.be.equal(comments);
+  });
 });

packages/babel-traverse/src/index.js

@@ -56,14 +56,14 @@ traverse.node = function (node: Object, opts: Object, scope: Object, state: Obje
   }
 };
 
-traverse.clearNode = function (node) {
-  t.removeProperties(node);
+traverse.clearNode = function (node, opts) {
+  t.removeProperties(node, opts);
 
   cache.path.delete(node);
 };
 
-traverse.removeProperties = function (tree) {
-  t.traverseFast(tree, traverse.clearNode);
+traverse.removeProperties = function (tree, opts) {
+  t.traverseFast(tree, traverse.clearNode, opts);
   return tree;
 };
 

packages/babel-types/src/index.js

@@ -455,40 +455,47 @@ toFastProperties(t.VISITOR_KEYS);
  * A prefix AST traversal implementation implementation.
  */
 
-export function traverseFast(node: Node, enter: (node: Node) => void) {
+export function traverseFast(node: Node, enter: (node: Node) => void, opts?: Object) {
   if (!node) return;
 
   let keys = t.VISITOR_KEYS[node.type];
   if (!keys) return;
 
-  enter(node);
+  opts = opts || {};
+  enter(node, opts);
 
   for (let key of keys) {
     let subNode = node[key];
 
     if (Array.isArray(subNode)) {
       for (let node of subNode) {
-        traverseFast(node, enter);
+        traverseFast(node, enter, opts);
       }
     } else {
-      traverseFast(subNode, enter);
+      traverseFast(subNode, enter, opts);
     }
   }
 }
 
-const CLEAR_KEYS: Array = t.COMMENT_KEYS.concat([
-  "tokens", "comments",
+const CLEAR_KEYS: Array = [
+  "tokens",
   "start", "end", "loc",
   "raw", "rawValue"
-]);
+];
+
+const CLEAR_KEYS_PLUS_COMMENTS: Array = t.COMMENT_KEYS.concat([
+  "comments"
+]).concat(CLEAR_KEYS);
 
 /**
  * Remove all of the _* properties from a node along with the additional metadata
  * properties like location data and raw token data.
  */
 
-export function removeProperties(node: Node): void {
-  for (let key of CLEAR_KEYS) {
+export function removeProperties(node: Node, opts?: Object): void {
+  opts = opts || {};
+  let map = opts.preserveComments ? CLEAR_KEYS : CLEAR_KEYS_PLUS_COMMENTS;
+  for (let key of map) {
     if (node[key] != null) node[key] = undefined;
   }
 
@@ -502,8 +509,8 @@ export function removeProperties(node: Node): void {
   }
 }
 
-export function removePropertiesDeep(tree: Node): Node {
-  traverseFast(tree, removeProperties);
+export function removePropertiesDeep(tree: Node, opts?: Object): Node {
+  traverseFast(tree, removeProperties, opts);
   return tree;
 }