dry up array comprehension for single block

README.md

@@ -149,9 +149,19 @@ browserify()
   .pipe(fs.createWriteStream("bundle.js"));
 ```
 
-## Caveats
+## Polyfill
 
-### Polyfill
+```javascript
+require("6to5/polyfill");
+```
+
+```javascript
+require("6to5").polyfill;
+```
+
+    $ 6to5-polyfill
+
+## Caveats
 
 ### Classes
 

lib/6to5/templates/array-comprehension-for-each.js

@@ -1,3 +1,3 @@
-OBJECT.forEach(function (KEY) {
+ARRAY.forEach(function (KEY) {
 
 });

lib/6to5/templates/array-comprehension-map-filter.js

@@ -0,0 +1,5 @@
+ARRAY.filter(function (KEY) {
+  return FILTER;
+}).map(function (KEY) {
+  return STATEMENT;
+});

lib/6to5/templates/array-comprehension-map.js

@@ -0,0 +1,3 @@
+ARRAY.map(function (KEY) {
+  return STATEMENT;
+});

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

@@ -1,6 +1,29 @@
 var util = require("../util");
+var _    = require("lodash");
 
 exports.ComprehensionExpression = function (node, parent, opts, generateUid) {
+  var blocks = node.blocks;
+
+  _.each(blocks, function (block) {
+    if (!block.of) {
+      throw util.errorWithNode(block, "for-in array comprehension is not supported");
+    }
+  });
+
+  if (blocks.length === 1) {
+    var block = blocks[0];
+
+    var templateName = "array-comprehension-map";
+    if (node.filter) templateName += "-filter";
+
+    return util.template(templateName, {
+      ARRAY: block.right,
+      KEY: block.left,
+      FILTER: node.filter,
+      STATEMENT: node.body
+    });
+  }
+
   var uid = generateUid("arr");
 
   var container = util.template("array-comprehension-container", {
@@ -13,13 +36,9 @@ exports.ComprehensionExpression = function (node, parent, opts, generateUid) {
   var returnStatement = body.pop();
 
   var build = function () {
-    var self = node.blocks.shift();
+    var self = blocks.shift();
     if (!self) return;
 
-    if (!self.of) {
-      throw util.errorWithNode(self, "for-in array comprehension is not supported");
-    }
-
     var child = build();
     if (!child) {
       // last item
@@ -29,6 +48,7 @@ exports.ComprehensionExpression = function (node, parent, opts, generateUid) {
         STATEMENT: node.body
       }, true);
 
+      // add a filter as this is our final stop
       if (node.filter) {
         var filter = util.template("if", {
           STATEMENT: node.filter
@@ -39,12 +59,10 @@ exports.ComprehensionExpression = function (node, parent, opts, generateUid) {
     }
 
     var container2 = util.template("array-comprehension-for-each", {
-      OBJECT: self.right,
-      KEY: self.left,
+      ARRAY: self.right,
+      KEY: self.left
     }, true);
-
     container2.expression.arguments[0].body.body = [child];
-
     return container2;
   };
 

lib/6to5/util.js

@@ -37,7 +37,7 @@ exports.parse = function (code, opts) {
 };
 
 exports.errorWithNode = function (node, msg) {
-  var err = new Error(msg);
+  var err = new SyntaxError(msg);
   err.lineNumber = node.loc.start.line;
   err.column = node.loc.start.column;
   return err;

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

@@ -1,12 +0,0 @@
-var seattlers = function () {
-  var _arr = [];
-  customers.forEach(function (c) {
-    if (c.city == "Seattle") {
-      _arr.push({
-        name: c.name,
-        age: c.age
-      });
-    }
-  });
-  return _arr;
-}();

test/fixtures/array-comprehension/multiple-if/actual.js

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

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

@@ -0,0 +1,14 @@
+var seattlers = function () {
+  var _arr = [];
+  countries.forEach(function (customers) {
+    customers.forEach(function (c) {
+      if (c.city == "Seattle") {
+        _arr.push({
+          name: c.name,
+          age: c.age
+        });
+      }
+    });
+  });
+  return _arr;
+}();

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

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

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

@@ -1,9 +1,3 @@
-var arr = (function () {
-  var _arr = [];
-
-  [1, 2, 3].forEach(function (i) {
-    _arr.push(i * i);
-  });
-
-  return _arr;
-})();
+var arr = [1, 2, 3].map(function (i) {
+  return i * i;
+});