Fix AMD define detection for arrow functions (#44)

Author: XhmikosR

index.js

@@ -90,6 +90,12 @@ function isModuleIdentifier(obj) {
   return obj.type && obj.type === 'Identifier' && obj.name === 'module';
 }
 
+function isFunctionLike(node) {
+  if (!node) return false;
+
+  return node.type === 'FunctionExpression' || node.type === 'ArrowFunctionExpression';
+}
+
 // module.exports.foo
 function isModuleExportsAttach(node) {
   if (!node.object || !node.object.object || !node.object.property) return false;
@@ -139,7 +145,7 @@ module.exports.isNamedForm = function(node) {
   return args && args.length === 3 &&
          (args[0].type === 'Literal' || args[0].type === 'StringLiteral') &&
          args[1].type === 'ArrayExpression' &&
-         args[2].type === 'FunctionExpression';
+         isFunctionLike(args[2]);
 };
 
 // define([deps], func)
@@ -150,7 +156,7 @@ module.exports.isDependencyForm = function(node) {
 
   return args && args.length === 2 &&
          args[0].type === 'ArrayExpression' &&
-         args[1].type === 'FunctionExpression';
+         isFunctionLike(args[1]);
 };
 
 // define(func(require))
@@ -162,7 +168,7 @@ module.exports.isFactoryForm = function(node) {
 
   // Node should have a function whose first param is 'require'
   return args && args.length === 1 &&
-         args[0].type === 'FunctionExpression' &&
+         isFunctionLike(args[0]) &&
          firstParamNode && firstParamNode.type === 'Identifier' &&
          firstParamNode.name === 'require';
 };
@@ -183,7 +189,7 @@ module.exports.isREMForm = function(node) {
   const args = node.arguments;
   const params = args.length > 0 ? args[0].params : null;
 
-  if (!args || args.length === 0 || args[0].type !== 'FunctionExpression' || params.length !== 3) {
+  if (!args || args.length === 0 || !isFunctionLike(args[0]) || params.length !== 3) {
     return false;
   }
 

test/is-define-amd.test.js

@@ -25,4 +25,15 @@ testSuite('detects a named form AMD define function call', () => {
   assert.not.ok(check('define();', types.isDefineAMD));
 });
 
+testSuite('detects AMD define with arrow callback returning class', () => {
+  assert.ok(check('define(["jquery"] , ($) => { class A {} return A;});', types.isDefineAMD));
+});
+
+testSuite('detects AMD forms with arrow callbacks', () => {
+  assert.ok(check('define("name", ["jquery"], ($) => $);', types.isDefineAMD));
+  assert.ok(check('define(["jquery"], ($) => $);', types.isDefineAMD));
+  assert.ok(check('define((require) => require("jquery"));', types.isDefineAMD));
+  assert.ok(check('define((require, exports, module) => { exports.x = true; });', types.isDefineAMD));
+});
+
 testSuite.run();