Support arrow functions in AMD wrappers

test/fixtures/amd.after.js

@@ -28,4 +28,34 @@ import moduleA from 'a';
 import moduleB from 'b';
 import moduleC from 'c';
 import 'd';
-console.log('xyz');
+console.log('xyz');
+console.log('c');
+export default {};
+export default {};
+console.log('last statement');
+console.log('d');
+import 'd';
+console.log('dont bother with d');
+import d from 'd';
+console.log('I need you D', d.isUsed());
+export default d.e;
+
+export default {
+	e: '123'
+};
+
+import 'e';
+
+export default {
+	e: 'e is ignored'
+};
+
+import e from 'e';
+
+export default {
+	e: `e is ${e.isUsed()}`
+};
+
+export default {
+	f: '123'
+};

test/fixtures/amd.before.js

@@ -41,3 +41,44 @@ define(['a', 'b', 'c', 'd'], function(moduleA, moduleB, moduleC) {
 	console.log('xyz');
 });
 
+define(() => {
+	console.log('c');
+});
+
+define(() => {
+	return {};
+});
+
+define(() => {
+  return {};
+  console.log('last statement');
+});
+
+define([], () => {
+	console.log('d');
+});
+
+define(['d'], () => {
+	console.log('dont bother with d');
+});
+
+define(['d'], (d) => {
+  console.log('I need you D', d.isUsed());
+  return d.e;
+});
+
+define([], () => ({
+	e: '123'
+}));
+
+define(['e'], () => ({
+	e: 'e is ignored'
+}));
+
+define(['e'], (e) => ({
+	e: `e is ${e.isUsed()}`
+}));
+
+define(() => ({
+	f: '123'
+}));

transforms/amd.js

@@ -9,6 +9,16 @@ module.exports = function(file, api, options) {
     var root = j(file.source);
     var leadingComment = root.find(j.Program).get('body', 0).node.leadingComments;
 
+    /**
+     * Get factory function whether arrow expression or block statement
+     * @param factory - Function body AST
+     */
+    function getFactoryBody(factory) {
+        return factory.body.type === 'BlockStatement'
+            ? factory.body.body
+            : [j.returnStatement(factory.body)]
+    }
+
     /**
      * Convert an `return` to `export default`.
      * @param body - Function body AST (Array)
@@ -39,7 +49,7 @@ module.exports = function(file, api, options) {
             if (p.value.arguments.length === 1) {
 
                 // convert `return` statement to `export default`
-                body = returnToExport(p.value.arguments[0].body.body);
+                body = returnToExport(getFactoryBody(p.value.arguments[0]));
 
                 return j(p.parent).replaceWith(body);
             }
@@ -55,7 +65,7 @@ module.exports = function(file, api, options) {
                 });
 
                 // add the body after the import statements
-                Array.prototype.push.apply(importStatements, p.value.arguments[1].body.body);
+                Array.prototype.push.apply(importStatements, getFactoryBody(p.value.arguments[1]));
 
                 // add any comments at the top
                 importStatements[0].comments = comments;