Basic React Hooks support (useImperativeHandle only)

README.md

@@ -368,7 +368,63 @@ we are getting this output:
 }
 ```
 
-### Types
+## React Hooks support
+
+If you are using React Hooks, react-docgen will now also find component methods defined directly via the `useImperativeHandle()` hook.
+
+> **Note**: react-docgen will not be able to grab the type definition if the type is imported or declared in a different file.
+
+### Example
+
+For the following component using `useImperativeHandle`:
+
+
+```js
+import React, { useImperativeHandle } from 'react';
+
+/**
+ * General component description.
+ */
+const MyComponent = React.forwardRef((props, ref) => {
+
+  useImperativeHandle(ref, () => ({
+    /**
+     * This is my method
+     */
+    myMethod: (arg1) => {},
+  }));
+
+  return /* ... */;
+});
+
+export default MyComponent;
+```
+
+we are getting this output:
+
+```json
+{
+  "description": "General component description.",
+  "displayName": "MyComponent",
+  "methods": [
+    {
+      "name": "myMethod",
+      "docblock": "This is my method",
+      "modifiers": [],
+      "params": [
+        {
+          "name": "arg1",
+          "optional": false
+        }
+      ],
+      "returns": null,
+      "description": "This is my method"
+    }
+  ]
+}
+```
+
+## Types
 
 Here is a list of all the available types and its result structure.
 

src/handlers/__tests__/componentMethodsHandler-test.js

@@ -8,7 +8,7 @@
 
 jest.mock('../../Documentation');
 
-import { parse } from '../../../tests/utils';
+import { parse, parseWithTemplate } from '../../../tests/utils';
 
 describe('componentMethodsHandler', () => {
   let documentation;
@@ -206,4 +206,97 @@ describe('componentMethodsHandler', () => {
       expect(documentation.methods).toMatchSnapshot();
     });
   });
+
+  describe('useImperativeHandle() methods', () => {
+    // We're not worried about doc-blocks here, simply about finding method(s)
+    // defined via the useImperativeHandle() hook.
+    const IMPERATIVE_TEMPLATE = [
+      'import React, { useImperativeHandle } from "react";',
+      '%s',
+    ].join('\n');
+
+    // To simplify the variations, each one ends up with the following in the
+    // parsed body:
+    //
+    // [0]: the react import
+    // [1]: the initial definition/declaration
+    // [2]: a React.forwardRef wrapper (or nothing)
+    //
+    // Note that in the cases where the React.forwardRef is used "inline" with
+    // the definition/declaration, there is no [2], and it will be skipped.
+
+    function testImperative(src) {
+      const parsed = parseWithTemplate(src, IMPERATIVE_TEMPLATE);
+      [1, 2].forEach(index => {
+        // reset the documentation, since we may test more than once!
+        documentation = new (require('../../Documentation'))();
+        const definition = parsed.get('body', index);
+        if (!definition.value) {
+          return;
+        }
+        componentMethodsHandler(documentation, definition);
+        expect(documentation.methods).toEqual([
+          {
+            docblock: null,
+            modifiers: [],
+            name: 'doFoo',
+            params: [],
+            returns: null,
+          },
+        ]);
+      });
+    }
+
+    it('finds inside a component in a variable declaration', () => {
+      testImperative(`
+      const Test = (props, ref) => {
+        useImperativeHandle(ref, () => ({
+          doFoo: ()=>{},
+        }));
+      };
+      React.forwardRef(Test);
+      `);
+    });
+
+    it('finds inside a component in an assignment', () => {
+      testImperative(`
+      Test = (props, ref) => {
+        useImperativeHandle(ref, () => ({
+          doFoo: ()=>{},
+        }));
+      };
+      `);
+    });
+
+    it('finds inside a function declaration', () => {
+      testImperative(`
+      function Test(props, ref) {
+        useImperativeHandle(ref, () => ({
+          doFoo: ()=>{},
+        }));
+      }
+      React.forwardRef(Test);
+      `);
+    });
+
+    it('finds inside an inlined React.forwardRef call with arrow function', () => {
+      testImperative(`
+      React.forwardRef((props, ref) => {
+        useImperativeHandle(ref, () => ({
+          doFoo: ()=>{},
+        }));
+      });
+      `);
+    });
+
+    it('finds inside an inlined React.forwardRef call with plain function', () => {
+      testImperative(`
+      React.forwardRef(function(props, ref) {
+        useImperativeHandle(ref, () => ({
+          doFoo: ()=>{},
+        }));
+      });
+      `);
+    });
+  });
 });

src/handlers/componentMethodsHandler.js

@@ -7,11 +7,13 @@
  * @flow
  */
 
-import { namedTypes as t } from 'ast-types';
+import { namedTypes as t, visit } from 'ast-types';
 import getMemberValuePath from '../utils/getMemberValuePath';
 import getMethodDocumentation from '../utils/getMethodDocumentation';
+import isReactBuiltinCall from '../utils/isReactBuiltinCall';
 import isReactComponentClass from '../utils/isReactComponentClass';
 import isReactComponentMethod from '../utils/isReactComponentMethod';
+import isReactForwardRefCall from '../utils/isReactForwardRefCall';
 import type Documentation from '../Documentation';
 import match from '../utils/match';
 import { traverseShallow } from '../utils/traverse';
@@ -65,6 +67,121 @@ function findAssignedMethods(scope, idPath) {
   return results;
 }
 
+// Finding the component itself depends heavily on how it's exported.
+// Conversely, finding any 'useImperativeHandle()' methods requires digging
+// through intervening assignments, declarations, and optionally a
+// React.forwardRef() call.
+function findUnderlyingComponentDefinition(exportPath) {
+  let path = exportPath;
+  let keepDigging = true;
+  let sawForwardRef = false;
+
+  // We can't use 'visit', because we're not necessarily climbing "down" the
+  // AST, we're following the logic flow *backwards* to the component
+  // definition. Once we do find what looks like the underlying functional
+  // component definition, *then* we can 'visit' downwards to find the call to
+  // useImperativeHandle, if it exists.
+  while (keepDigging && path) {
+    // Using resolveToValue automatically gets the "value" from things like
+    // assignments or identifier references.  Putting this here removes the need
+    // to call it in a bunch of places on a per-type basis.
+    path = resolveToValue(path);
+    const node = path.node;
+
+    // Rather than using ast-types 'namedTypes' (t) checks, we 'switch' on the
+    // `node.type` value.  We lose the "is a" aspect (like a CallExpression "is
+    // a(n)" Expression), but our handling very much depends on the *exact* node
+    // type, so that's an acceptable compromise.
+    switch (node.type) {
+      case 'VariableDeclaration':
+        path = path.get('declarations');
+        if (path.value && path.value.length === 1) {
+          path = path.get(0);
+        } else {
+          path = null;
+        }
+        break;
+
+      case 'ExpressionStatement':
+        path = path.get('expression');
+        break;
+
+      case 'CallExpression':
+        // FUTURE: Can we detect other common HOCs that we could drill through?
+        if (isReactForwardRefCall(path) && !sawForwardRef) {
+          sawForwardRef = true;
+          path = path.get('arguments', 0);
+        } else {
+          path = null;
+        }
+        break;
+
+      case 'ArrowFunctionExpression':
+      case 'FunctionDeclaration':
+      case 'FunctionExpression':
+        // get the body and visit for useImperativeHandle!
+        path = path.get('body');
+        keepDigging = false;
+        break;
+
+      default:
+        // Any other type causes us to bail.
+        path = null;
+    }
+  }
+
+  return path;
+}
+
+function findImperativeHandleMethods(exportPath) {
+  const path = findUnderlyingComponentDefinition(exportPath);
+
+  if (!path) {
+    return [];
+  }
+
+  const results = [];
+  visit(path, {
+    visitCallExpression: function(callPath) {
+      // We're trying to handle calls to React's useImperativeHandle.  If this
+      // isn't, we can stop visiting this node path immediately.
+      if (!isReactBuiltinCall(callPath, 'useImperativeHandle')) {
+        return false;
+      }
+
+      // The standard use (and documented example) is:
+      //
+      //   useImperativeHandle(ref, () => ({ name: () => {}, ...}))
+      //                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+      //
+      // ... so we only handle a second argument (index 1) that is an
+      // ArrowFunctionExpression and whose body is an ObjectExpression.
+      const arg = callPath.get('arguments', 1);
+
+      if (!t.ArrowFunctionExpression.check(arg.node)) {
+        return false;
+      }
+
+      const body = arg.get('body');
+      if (!t.ObjectExpression.check(body.node)) {
+        return false;
+      }
+
+      // We found the object body, now add all of the properties as methods.
+      traverseShallow(body.get('properties'), {
+        visitProperty: prop => {
+          results.push(prop);
+          return false;
+        },
+      });
+
+      return false;
+    },
+  });
+
+  return results;
+}
+
 /**
  * Extract all flow types for the methods of a react component. Doesn't
  * return any react specific lifecycle methods.
@@ -109,6 +226,12 @@ export default function componentMethodsHandler(
     methodPaths = findAssignedMethods(path.parent.scope, path.get('id'));
   }
 
+  // Also look for any methods that come from useImperativeHandle() calls.
+  const impMethodPaths = findImperativeHandleMethods(path);
+  if (impMethodPaths && impMethodPaths.length > 0) {
+    methodPaths = methodPaths.concat(impMethodPaths);
+  }
+
   documentation.set(
     'methods',
     methodPaths.map(getMethodDocumentation).filter(Boolean),