Adding support in for module pattern, adding Function Expressions to …

…the current visitors added for pure components.

src/index.js

@@ -95,6 +95,9 @@ export default function ({ Plugin, types: t }) {
     if (node.id) {
       return node.id.name;
     }
+    if (t.isFunctionExpression && t.isProperty(parent) && parent.key) {
+      return parent.key.name;
+    }
     if (t.isExportDefaultDeclaration(parent) &&
          (t.isArrowFunctionExpression(node) || t.isFunctionExpression(node))) {
       return file.opts.basename;
@@ -297,7 +300,7 @@ export default function ({ Plugin, types: t }) {
         }
       },
 
-      'FunctionDeclaration|ArrowFunctionExpression': {
+      'FunctionDeclaration|FunctionExpression|ArrowFunctionExpression': {
         enter(node, parent, scope, file) {
           if (!this.state[depthKey]) {
             this.state[depthKey] = 0;
@@ -312,16 +315,27 @@ export default function ({ Plugin, types: t }) {
             return;
           }
 
+          // No-op for `render` methods in this visitor
+          if (parent.key && parent.key.name === 'render') {
+            return;
+          }
+
           delete this.state[foundJSXKey];
 
           const wrapReactComponentId = this.state[wrapComponentIdKey];
           const uniqueId = addComponentRecord(node, parent, scope, file, this.state);
           const bindingId = node.id;
 
-          if (t.isArrowFunctionExpression(node) &&
-             (t.isVariableDeclarator(parent) ||
-                t.isExportDefaultDeclaration(parent))
-          ) {
+          if (t.isArrowFunctionExpression(node)) {
+            if (t.isVariableDeclarator(parent) || t.isExportDefaultDeclaration(parent)) {
+              return t.callExpression(
+                t.callExpression(wrapReactComponentId, [t.literal(uniqueId)]),
+                [node]
+              );
+            }
+          }
+
+          if (t.isFunctionExpression(node)) {
             return t.callExpression(
               t.callExpression(wrapReactComponentId, [t.literal(uniqueId)]),
               [node]

test/fixtures/pure-function/actual.js

@@ -1,6 +1,6 @@
 import React from 'react';
 
-function DefaultFunction() {
+function FunctionDeclaration() {
   return <div />;
 }
 
@@ -12,6 +12,15 @@ function DefaultExportFunction() {
   return <div />;
 }
 
+const ModulePattern = {
+  NoProps() {
+    return <div />;
+  },
+  WithProps(props) {
+    return <div {...props} />;
+  }
+};
+
 export { NamedExportFunction };
 export default DefaultExportFunction;
 export default function () {

test/fixtures/pure-function/expected.js

@@ -17,8 +17,8 @@ Object.defineProperty(exports, '__esModule', {
 var _react2 = _interopRequireDefault(_react);
 
 var _components = {
-  _$DefaultFunction: {
-    displayName: 'DefaultFunction',
+  _$FunctionDeclaration: {
+    displayName: 'FunctionDeclaration',
     isFunction: true
   },
   _$NamedExportFunction: {
@@ -29,6 +29,14 @@ var _components = {
     displayName: 'DefaultExportFunction',
     isFunction: true
   },
+  _$NoProps: {
+    displayName: 'NoProps',
+    isFunction: true
+  },
+  _$WithProps: {
+    displayName: 'WithProps',
+    isFunction: true
+  },
   _$actual: {
     displayName: 'actual',
     isFunction: true
@@ -57,10 +65,11 @@ function _wrapComponent(uniqueId) {
 
 function _interopRequireDefault(obj) { return obj && obj.__esModule ? obj : { 'default': obj }; }
 
-function DefaultFunction() {
+function FunctionDeclaration() {
   return _react2['default'].createElement('div', null);
 }
 
+FunctionDeclaration = _wrapComponent('_$FunctionDeclaration')(FunctionDeclaration)
 function NamedExportFunction() {
   return _react2['default'].createElement('div', null);
 }
@@ -69,8 +78,16 @@ exports.NamedExportFunction = NamedExportFunction = _wrapComponent('_$NamedExpor
 function DefaultExportFunction() {
   return _react2['default'].createElement('div', null);
 }
-
 DefaultExportFunction = _wrapComponent('_$DefaultExportFunction')(DefaultExportFunction)
+var ModulePattern = {
+  NoProps: _wrapComponent('_$NoProps')(function () {
+    return _react2['default'].createElement('div', null);
+  }),
+  WithProps: _wrapComponent('_$WithProps')(function (props) {
+    return _react2['default'].createElement('div', props);
+  })
+};
+
 exports.NamedExportFunction = NamedExportFunction;
 exports['default'] = DefaultExportFunction;
 exports['default'] = _wrapComponent('_$actual')(function () {