make array comprehensions use for-of unless using an array literal - fixes #98

lib/6to5/transformers/array-comprehension.js

@@ -1,11 +1,11 @@
 var util = require("../util");
-var b    = require("../builders");
+var t    = require("../types");
 
-var single = function (node) {
+var singleArrayExpression = function (node) {
   var block = node.blocks[0];
 
-  var templateName = "array-comprehension-map";
-  if (node.filter) templateName = "array-comprehension-filter";
+  var templateName = "array-expression-comprehension-map";
+  if (node.filter) templateName = "array-expression-comprehension-filter";
 
   var result = util.template(templateName, {
     STATEMENT: node.body,
@@ -45,19 +45,15 @@ var multiple = function (node, file) {
 
       // add a filter as this is our final stop
       if (node.filter) {
-        child = b.ifStatement(node.filter, b.blockStatement([child]));
+        child = t.ifStatement(node.filter, t.blockStatement([child]));
       }
     }
 
-    var container2 = util.template("array-comprehension-for-each", {
-      ARRAY: self.right,
-      KEY:   self.left
-    }, true);
-
-    // set function body
-    container2.expression.arguments[0].body.body = [child];
-
-    return container2;
+    return t.forOfStatement(
+      t.variableDeclaration("var", [t.variableDeclarator(self.left)]),
+      self.right,
+      t.blockStatement([child])
+    );
   };
 
   body.push(build());
@@ -67,8 +63,8 @@ var multiple = function (node, file) {
 };
 
 exports.ComprehensionExpression = function (node, parent, file) {
-  if (node.blocks.length === 1) {
-    return single(node);
+  if (node.blocks.length === 1 && t.isArrayExpression(node.blocks[0].right)) {
+    return singleArrayExpression(node);
   } else {
     return multiple(node, file);
   }

test/fixtures/transformation/array-comprehension/array-expression-single-if/actual.js

@@ -0,0 +1 @@
+var arr = [for (i of [1, 2, 3]) if (i > 1) i * i];

test/fixtures/transformation/array-comprehension/array-expression-single-if/expected.js

@@ -0,0 +1,7 @@
+"use strict";
+
+var arr = [1, 2, 3].filter(function (i) {
+  return i > 1;
+}).map(function (i) {
+  return i * i;
+});

test/fixtures/transformation/array-comprehension/multiple-if/expected.js

@@ -2,16 +2,21 @@
 
 var seattlers = (function () {
   var _arr = [];
-  countries.forEach(function (customers) {
-    customers.forEach(function (c) {
+
+  for (var _iterator = countries[Symbol.iterator](), _step; !(_step = _iterator.next()).done;) {
+    var customers = _step.value;
+
+    for (var _iterator2 = customers[Symbol.iterator](), _step2; !(_step2 = _iterator2.next()).done;) {
+      var c = _step2.value;
+
       if (c.city == "Seattle") {
         _arr.push({
           name: c.name,
           age: c.age
         });
       }
-    });
-  });
+    }
+  }
 
   return _arr;
-})();
+})();

test/fixtures/transformation/array-comprehension/multiple/expected.js

@@ -2,11 +2,16 @@
 
 var arr = (function () {
   var _arr = [];
-  "abcdefgh".split("").forEach(function (x) {
-    "12345678".split("").forEach(function (y) {
+
+  for (var _iterator = "abcdefgh".split("")[Symbol.iterator](), _step; !(_step = _iterator.next()).done;) {
+    var x = _step.value;
+
+    for (var _iterator2 = "12345678".split("")[Symbol.iterator](), _step2; !(_step2 = _iterator2.next()).done;) {
+      var y = _step2.value;
+
       _arr.push(x + y);
-    });
-  });
+    }
+  }
 
   return _arr;
-})();
+})();

test/fixtures/transformation/array-comprehension/single-if/actual.js

@@ -1 +0,0 @@
-var seattlers = [for (c of customers) if (c.city == "Seattle") { name: c.name, age: c.age }];

test/fixtures/transformation/array-comprehension/single-if/expected.js

@@ -1,10 +0,0 @@
-"use strict";
-
-var seattlers = customers.filter(function (c) {
-  return c.city == "Seattle";
-}).map(function (c) {
-  return {
-    name: c.name,
-    age: c.age
-  };
-});