diff --git a/src/index.js b/src/index.js
index 6f35b9b..95f963e 100644
--- a/src/index.js
+++ b/src/index.js
@@ -1,51 +1,68 @@
 import * as _ from 'lodash';
 import * as reactDocs from 'react-docgen';
 import isReactComponentClass from './isReactComponentClass';
+import isStatelessComponent from './isStatelessComponent';
 import * as p from 'path';
 
-export default function({types: t}) {
+export default function ({types: t}) {
   return {
     visitor: {
       Class(path, state) {
-        const {
-          node,
-          scope,
-        } = path;
-
-        if(isReactComponentClass(path)){
-          injectReactDocgenInfo(path, this.file.code, t);
-          injectDocgenGlobal(path, state, t);
+        if(isReactComponentClass(path)) {
+          const className = path.node.id.name;
+          injectReactDocgenInfo(className, path, state, this.file.code, t);
         }
       },
-      'FunctionDeclaration|FunctionExpression|ArrowFunctionExpression'(path) {
+      'FunctionDeclaration|FunctionExpression|ArrowFunctionExpression'(path, state) {
+        if(isStatelessComponent(path)) {
+          const className = path.parentPath.node.id.name;
+          injectReactDocgenInfo(className, path, state, this.file.code, t);
+        }
       },
     }
   };
 }
 
+function alreadyVisited(program, t) {
+  return program.node.body.some(node => {
+    if(t.isExpressionStatement(node) &&
+       t.isAssignmentExpression(node.expression) &&
+       t.isMemberExpression(node.expression.left)
+      ) {
+      return node.expression.left.property.name === '__docgenInfo';
+    }
+    return false;
+  });
+}
+
 
-function injectReactDocgenInfo(path, code, t) {
-  if(!t.isProgram(path.parentPath.node)) {
+function injectReactDocgenInfo(className, path, state, code, t) {
+  const program = path.scope.getProgramParent().path;
+
+  if(alreadyVisited(program, t)) {
     return;
   }
+
   const docObj = reactDocs.parse(code);
   const docNode = buildObjectExpression(docObj, t);
   const docgenInfo = t.expressionStatement(
     t.assignmentExpression(
       "=",
-      t.memberExpression(t.identifier(path.node.id.name), t.identifier('__docgenInfo')),
+      t.memberExpression(t.identifier(className), t.identifier('__docgenInfo')),
       docNode
     ));
-  path.parentPath.pushContainer('body', docgenInfo);
-
+  program.pushContainer('body', docgenInfo);
+  injectDocgenGlobal(className, path, state, t);
 }
 
-function injectDocgenGlobal(path, state, t) {
-  if(!state.opts.DOC_GEN_GLOBAL || !t.isProgram(path.parentPath.node)) {
+function injectDocgenGlobal(className, path, state, t) {
+  const program = path.scope.getProgramParent().path;
+
+  if(!state.opts.DOC_GEN_GLOBAL) {
     return;
   }
+
   const globalName = state.opts.DOC_GEN_GLOBAL;
-  const className = path.node.id.name;
   const filePath = p.relative('./', p.resolve('./', path.hub.file.opts.filename));
   const globalNode = t.ifStatement(
     t.binaryExpression(
@@ -86,7 +103,7 @@ function injectDocgenGlobal(path, state, t) {
       )
     ])
   );
-  path.parentPath.pushContainer('body', globalNode);
+  program.pushContainer('body', globalNode);
 }
 
 function buildObjectExpression(obj, t){
@@ -109,7 +126,7 @@ function buildObjectExpression(obj, t){
     return t.numericLiteral(obj);
   } else if (_.isArray(obj)) {
     const children = [];
-    obj.forEach(function(val) {
+    obj.forEach(function (val) {
       children.push(buildObjectExpression(val, t));
     });
     return t.ArrayExpression(children);
diff --git a/src/isStatelessComponent.js b/src/isStatelessComponent.js
new file mode 100644
index 0000000..5d950ba
--- /dev/null
+++ b/src/isStatelessComponent.js
@@ -0,0 +1,91 @@
+function isJSXElementOrReactCreateElement(node) {
+  const {
+    type,
+    callee,
+  } = node;
+
+  if (type === 'JSXElement') {
+    return true;
+  }
+
+  if (callee && callee.object && callee.object.name === 'React' &&
+    callee.property.name === 'createElement') {
+    return true;
+  }
+
+  return false;
+}
+
+function isReturningJSXElement(path) {
+  /**
+   * Early exit for ArrowFunctionExpressions, there is no ReturnStatement node.
+   */
+  if (path.node.init && path.node.init.body && isJSXElementOrReactCreateElement(path.node.init.body)) {
+    return true;
+  }
+
+  let visited = false;
+
+  path.traverse({
+    ReturnStatement(path2) {
+      // We have already found what we are looking for.
+      if (visited) {
+        return;
+      }
+
+      const argument = path2.get('argument');
+
+      // Nothing is returned
+      if (!argument.node) {
+        return;
+      }
+
+      if (isJSXElementOrReactCreateElement(argument.node)) {
+        visited = true;
+        return;
+      }
+
+      if (argument.node.type === 'CallExpression') {
+        const name = argument.get('callee').node.name;
+        const binding = path.scope.getBinding(name);
+
+        if (!binding) {
+          return;
+        }
+
+        if (isReturningJSXElement(binding.path)) {
+          visited = true;
+        }
+      }
+    },
+  });
+
+  return visited;
+}
+
+const validPossibleStatelessComponentTypes = [
+  'Property',
+  'VariableDeclarator',
+  'FunctionDeclaration',
+  'ArrowFunctionExpression',
+];
+
+/**
+ * Returns `true` if the path represents a function which returns a JSXElement
+ */
+export default function isStatelessComponent(path) {
+  const node = path.node;
+
+  if (validPossibleStatelessComponentTypes.indexOf(node.type) === -1) {
+    return false;
+  }
+
+  if(path.get('body').get('type').node == 'JSXElement') {
+    return true;
+  }
+  if (isReturningJSXElement(path)) {
+    return true;
+  }
+
+  return false;
+}
diff --git a/test/fixtures/case3/actual.js b/test/fixtures/case3/actual.js
new file mode 100644
index 0000000..92f7170
--- /dev/null
+++ b/test/fixtures/case3/actual.js
@@ -0,0 +1,27 @@
+import React from 'react';
+
+const Button = ({ children, onClick, style = {} }) => (
+  <button
+    style={{ }}
+    onClick={onClick}
+  >
+    {children}
+  </button>
+);
+
+Button.propTypes = {
+  children: React.PropTypes.string.isRequired,
+  onClick: React.PropTypes.func,
+  style: React.PropTypes.object,
+};
+
+export default Button;
+
+let A;
+A = [1,2,2,2];
+
+function abc() {
+  let c = function cef() {
+    A = 'str';
+  };
+}
diff --git a/test/fixtures/case3/expected.js b/test/fixtures/case3/expected.js
new file mode 100644
index 0000000..432e988
--- /dev/null
+++ b/test/fixtures/case3/expected.js
@@ -0,0 +1,79 @@
+'use strict';
+
+Object.defineProperty(exports, "__esModule", {
+  value: true
+});
+
+var _react = require('react');
+
+var _react2 = _interopRequireDefault(_react);
+
+function _interopRequireDefault(obj) { return obj && obj.__esModule ? obj : { default: obj }; }
+
+var Button = function Button(_ref) {
+  var children = _ref.children;
+  var onClick = _ref.onClick;
+  var _ref$style = _ref.style;
+  var style = _ref$style === undefined ? {} : _ref$style;
+  return _react2.default.createElement(
+    'button',
+    {
+      style: {},
+      onClick: onClick
+    },
+    children
+  );
+};
+
+Button.propTypes = {
+  children: _react2.default.PropTypes.string.isRequired,
+  onClick: _react2.default.PropTypes.func,
+  style: _react2.default.PropTypes.object
+};
+
+exports.default = Button;
+
+
+var A = void 0;
+A = [1, 2, 2, 2];
+
+function abc() {
+  var c = function cef() {
+    A = 'str';
+  };
+}
+Button.__docgenInfo = {
+  description: '',
+  methods: [],
+  props: {
+    children: {
+      type: {
+        name: 'string'
+      },
+      required: true,
+      description: ''
+    },
+    onClick: {
+      type: {
+        name: 'func'
+      },
+      required: false,
+      description: ''
+    },
+    style: {
+      type: {
+        name: 'object'
+      },
+      required: false,
+      description: ''
+    }
+  }
+};
+
+if (typeof STORYBOOK_REACT_CLASSES !== 'undefined') {
+  STORYBOOK_REACT_CLASSES['test/fixtures/case3/actual.js'] = {
+    name: 'Button',
+    docgenInfo: Button.__docgenInfo,
+    path: 'test/fixtures/case3/actual.js'
+  };
+}
diff --git a/test/fixtures/case4/actual.js b/test/fixtures/case4/actual.js
new file mode 100644
index 0000000..25c7430
--- /dev/null
+++ b/test/fixtures/case4/actual.js
@@ -0,0 +1,29 @@
+import React from 'react';
+
+const Button = ({ children, onClick, style = {} }) => {
+  return (
+      <button
+    style={{ }}
+    onClick={onClick}
+      >
+      {children}
+    </button>
+  );
+};
+
+Button.propTypes = {
+  children: React.PropTypes.string.isRequired,
+  onClick: React.PropTypes.func,
+  style: React.PropTypes.object,
+};
+
+export default Button;
+
+let A;
+A = [1,2,2,2];
+
+function abc() {
+  let c = function cef() {
+    A = 'str';
+  };
+}
diff --git a/test/fixtures/case4/expected.js b/test/fixtures/case4/expected.js
new file mode 100644
index 0000000..615d77e
--- /dev/null
+++ b/test/fixtures/case4/expected.js
@@ -0,0 +1,80 @@
+'use strict';
+
+Object.defineProperty(exports, "__esModule", {
+  value: true
+});
+
+var _react = require('react');
+
+var _react2 = _interopRequireDefault(_react);
+
+function _interopRequireDefault(obj) { return obj && obj.__esModule ? obj : { default: obj }; }
+
+var Button = function Button(_ref) {
+  var children = _ref.children;
+  var onClick = _ref.onClick;
+  var _ref$style = _ref.style;
+  var style = _ref$style === undefined ? {} : _ref$style;
+
+  return _react2.default.createElement(
+    'button',
+    {
+      style: {},
+      onClick: onClick
+    },
+    children
+  );
+};
+
+Button.propTypes = {
+  children: _react2.default.PropTypes.string.isRequired,
+  onClick: _react2.default.PropTypes.func,
+  style: _react2.default.PropTypes.object
+};
+
+exports.default = Button;
+
+
+var A = void 0;
+A = [1, 2, 2, 2];
+
+function abc() {
+  var c = function cef() {
+    A = 'str';
+  };
+}
+Button.__docgenInfo = {
+  description: '',
+  methods: [],
+  props: {
+    children: {
+      type: {
+        name: 'string'
+      },
+      required: true,
+      description: ''
+    },
+    onClick: {
+      type: {
+        name: 'func'
+      },
+      required: false,
+      description: ''
+    },
+    style: {
+      type: {
+        name: 'object'
+      },
+      required: false,
+      description: ''
+    }
+  }
+};
+
+if (typeof STORYBOOK_REACT_CLASSES !== 'undefined') {
+  STORYBOOK_REACT_CLASSES['test/fixtures/case4/actual.js'] = {
+    name: 'Button',
+    docgenInfo: Button.__docgenInfo,
+    path: 'test/fixtures/case4/actual.js'
+  };
+}
diff --git a/test/fixtures/example/.babelrc b/test/fixtures/example/.babelrc
deleted file mode 100644
index e624888..0000000
--- a/test/fixtures/example/.babelrc
+++ /dev/null
@@ -1,5 +0,0 @@
-{
-  "plugins": [
-    ["../../../src"]
-  ]
-}
diff --git a/test/index.js b/test/index.js
index 7bc1a6e..a476edf 100644
--- a/test/index.js
+++ b/test/index.js
@@ -26,12 +26,12 @@ describe('Add propType doc to react classes', () => {
         ],
         babelrc: false
       };
+
       const actual = transformFileSync(actualPath, options).code;
       //fs.writeFileSync(path.join(fixtureDir, 'actual-output.js'), actual);
       const expected = fs.readFileSync(
           path.join(fixtureDir, 'expected.js')
       ).toString();
-
       assert.equal(trim(actual), trim(expected));
     });
   });