ensureBlock on multiple node replacements - fixes #229

lib/6to5/traverse/index.js

@@ -50,6 +50,10 @@ function traverse(parent, callbacks, opts) {
         if (result != null) {
           updated = true;
           node = obj[key] = result;
+
+          if (_.contains(t.STATEMENT_OR_BLOCK_KEYS, key) && !t.isBlockStatement(obj)) {
+            t.ensureBlock(obj, key);
+          }
         }
       };
 

lib/6to5/types/index.js

@@ -12,6 +12,8 @@ var addAssert = function (type, is) {
   };
 };
 
+t.STATEMENT_OR_BLOCK_KEYS = ["consequent", "body"];
+
 //
 
 t.VISITOR_KEYS = require("./visitor-keys");
@@ -167,8 +169,9 @@ t.isValidIdentifier = function (name) {
   return _.isString(name) && esutils.keyword.isIdentifierName(name) && !esutils.keyword.isKeywordES6(name, true);
 };
 
-t.ensureBlock = function (node) {
-  node.body = t.toBlock(node.body, node);
+t.ensureBlock = function (node, key) {
+  key = key || "body";
+  node[key] = t.toBlock(node[key], node);
 };
 
 t.toStatement = function (node, ignore) {