cleanup & handle more cases

src/index.js

@@ -66,16 +66,10 @@ function injectReactDocgenInfo(path, state, code, t) {
         docNode
       ));
 
-    const defaultExportDeclaration = program.get('body').find((node) => {
-      const isHoC = t.isCallExpression(node.get('declaration').node)
-      return (
-        t.isExportDefaultDeclaration(node) &&
-        (node.get('declaration').node.name === exportName || isHoC)
-      );
-    });
+    const exportPath = program.get('body').find((node) => isExportCurrent(node, exportName, t));
 
-    if (defaultExportDeclaration) {
-      defaultExportDeclaration.insertBefore(docgenInfo);
+    if (exportPath) {
+      exportPath.insertBefore(docgenInfo);
     } else {
       program.pushContainer('body', docgenInfo);
     }
@@ -164,3 +158,34 @@ function buildObjectExpression(obj, t){
     return t.nullLiteral();
   }
 }
+
+function getComponentFromHoC(path) {
+  if (path.isCallExpression()) {
+    return getComponentFromHoC(path.get('arguments.0'));
+  }
+  return path.isIdentifier() ? path.node.name : null;
+}
+
+function isExportCurrent(path, exportName, t) {
+  if (t.isExportDefaultDeclaration(path)) {
+    const decl = path.get('declaration');
+
+    const identifier = (
+      decl.isIdentifier() // export default MyComp
+        ? decl.node.name
+        : getComponentFromHoC(decl) // export default withHoC(MyComp)
+    );
+
+    if (identifier === exportName) {
+      return true
+    }
+  }
+
+  if (t.isExportNamedDeclaration(path)) {
+    return path.get('specifiers').find((sp) => (
+      sp.node.exported.name === exportName
+    ))
+  }
+
+  return false
+}

test/fixtures/hoc-multiple/actual.js

@@ -17,3 +17,14 @@ Component.propTypes = {
 }
 
 export default withHoc()(deeperHoc(Component))
+
+class CompA extends React.Component {
+  render() { return null }
+}
+
+CompA.propTypes = {
+  /** Fancy styles in here */
+  myProp: React.PropTypes.object,
+}
+
+export { CompA }

test/fixtures/hoc-multiple/expected.js

@@ -3,7 +3,7 @@
 Object.defineProperty(exports, "__esModule", {
   value: true
 });
-exports["default"] = void 0;
+exports.CompA = exports["default"] = void 0;
 
 function _typeof(obj) { if (typeof Symbol === "function" && typeof Symbol.iterator === "symbol") { _typeof = function _typeof(obj) { return typeof obj; }; } else { _typeof = function _typeof(obj) { return obj && typeof Symbol === "function" && obj.constructor === Symbol && obj !== Symbol.prototype ? "symbol" : typeof obj; }; } return _typeof(obj); }
 
@@ -95,10 +95,60 @@ var _default = withHoc()(deeperHoc(Component));
 
 exports["default"] = _default;
 
+var CompA =
+/*#__PURE__*/
+function (_React$Component2) {
+  _inherits(CompA, _React$Component2);
+
+  function CompA() {
+    _classCallCheck(this, CompA);
+
+    return _possibleConstructorReturn(this, _getPrototypeOf(CompA).apply(this, arguments));
+  }
+
+  _createClass(CompA, [{
+    key: "render",
+    value: function render() {
+      return null;
+    }
+  }]);
+
+  return CompA;
+}(React.Component);
+
+exports.CompA = CompA;
+CompA.propTypes = {
+  /** Fancy styles in here */
+  myProp: React.PropTypes.object
+};
+CompA.__docgenInfo = {
+  "description": "",
+  "methods": [],
+  "displayName": "CompA",
+  "props": {
+    "myProp": {
+      "type": {
+        "name": "custom",
+        "raw": "React.PropTypes.object"
+      },
+      "required": false,
+      "description": "Fancy styles in here"
+    }
+  }
+};
+
 if (typeof STORYBOOK_REACT_CLASSES !== "undefined") {
   STORYBOOK_REACT_CLASSES["test/fixtures/hoc-multiple/actual.js"] = {
     name: "Component",
     docgenInfo: Component.__docgenInfo,
     path: "test/fixtures/hoc-multiple/actual.js"
   };
 }
+
+if (typeof STORYBOOK_REACT_CLASSES !== "undefined") {
+  STORYBOOK_REACT_CLASSES["test/fixtures/hoc-multiple/actual.js"] = {
+    name: "CompA",
+    docgenInfo: CompA.__docgenInfo,
+    path: "test/fixtures/hoc-multiple/actual.js"
+  };
+}

test/fixtures/multiple-exports/expected.js

@@ -101,15 +101,6 @@ exports.ErrorBox2 = ErrorBox2;
 ErrorBox2.propTypes = {
   children2: _react["default"].PropTypes.node.isRequired
 };
-
-if (typeof STORYBOOK_REACT_CLASSES !== "undefined") {
-  STORYBOOK_REACT_CLASSES["test/fixtures/multiple-exports/actual.js"] = {
-    name: "ErrorBox",
-    docgenInfo: ErrorBox.__docgenInfo,
-    path: "test/fixtures/multiple-exports/actual.js"
-  };
-}
-
 ErrorBox2.__docgenInfo = {
   "description": "",
   "methods": [],
@@ -125,6 +116,14 @@ ErrorBox2.__docgenInfo = {
   }
 };
 
+if (typeof STORYBOOK_REACT_CLASSES !== "undefined") {
+  STORYBOOK_REACT_CLASSES["test/fixtures/multiple-exports/actual.js"] = {
+    name: "ErrorBox",
+    docgenInfo: ErrorBox.__docgenInfo,
+    path: "test/fixtures/multiple-exports/actual.js"
+  };
+}
+
 if (typeof STORYBOOK_REACT_CLASSES !== "undefined") {
   STORYBOOK_REACT_CLASSES["test/fixtures/multiple-exports/actual.js"] = {
     name: "ErrorBox2",

test/index.js

@@ -8,7 +8,7 @@ function trim(str) {
   return str.replace(/^\s+|\s+$/, '');
 }
 
-describe('Add propType doc to react classes', () => {
+describe('Add propType doc to react components', () => {
   const fixturesDir = path.join(__dirname, 'fixtures');
   fs.readdirSync(fixturesDir).map((caseName) => {
     // Ignore macOS directory files