useImperativeHandle support

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.
 

packages/react-docgen/src/handlers/__tests__/componentMethodsHandler-test.ts

@@ -7,8 +7,10 @@ import type {
   ClassDeclaration,
   ExportDefaultDeclaration,
   FunctionDeclaration,
+  VariableDeclaration,
 } from '@babel/types';
 import type { NodePath } from '@babel/traverse';
+import type { ComponentNode } from '../../resolver';
 
 jest.mock('../../Documentation');
 
@@ -413,4 +415,115 @@ 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.
+    // To simplify the variations, each one ends up with the following in the
+    // parsed body:
+    //
+    // [0] : the initial definition/declaration
+    // [1] : a React.forwardRef wrapper (or nothing)
+    // [last]: the react import
+    //
+    // Note that in the cases where the React.forwardRef is used "inline" with
+    // the definition/declaration, there is no [1], and it will be skipped.
+    function testImperative(src, paths: Array<string | null> = [null]) {
+      const srcWithImport = `
+        ${src}
+        import React, { useImperativeHandle } from "react";
+      `;
+
+      paths.forEach((path, index) => {
+        const parsed = parse.statement<VariableDeclaration>(
+          srcWithImport,
+          index,
+        );
+        const componentDefinition =
+          path != null
+            ? (parsed.get(path) as NodePath<ComponentNode>)
+            : (parsed as unknown as NodePath<ComponentNode>);
+
+        // reset the documentation, since we may test more than once!
+        documentation = new Documentation() as Documentation & DocumentationMock;
+        componentMethodsHandler(documentation, componentDefinition);
+        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);
+      `,
+        ['declarations.0.init', null],
+      );
+    });
+
+    it.only('finds inside a component in an assignment', () => {
+      testImperative(
+        `
+      Test = (props, ref) => {
+        useImperativeHandle(ref, () => ({
+          doFoo: ()=>{},
+        }));
+      };
+      `,
+        ['expression.right'],
+      );
+    });
+
+    it('finds inside a function declaration', () => {
+      testImperative(
+        `
+      function Test(props, ref) {
+        useImperativeHandle(ref, () => ({
+          doFoo: ()=>{},
+        }));
+      }
+      React.forwardRef(Test);
+      `,
+        [null, null],
+      );
+    });
+
+    it('finds inside an inlined React.forwardRef call with arrow function', () => {
+      testImperative(
+        `
+      React.forwardRef((props, ref) => {
+        useImperativeHandle(ref, () => ({
+          doFoo: ()=>{},
+        }));
+      });
+      `,
+        [null],
+      );
+    });
+
+    it('finds inside an inlined React.forwardRef call with plain function', () => {
+      testImperative(
+        `
+      React.forwardRef(function(props, ref) {
+        useImperativeHandle(ref, () => ({
+          doFoo: ()=>{},
+        }));
+      });
+      `,
+        [null],
+      );
+    });
+  });
 });

packages/react-docgen/src/handlers/componentMethodsHandler.ts

@@ -1,14 +1,16 @@
 import getMemberValuePath from '../utils/getMemberValuePath';
 import type { MethodNodePath } from '../utils/getMethodDocumentation';
 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 { shallowIgnoreVisitors } from '../utils/traverse';
 import resolveToValue from '../utils/resolveToValue';
 import type { NodePath, Scope } from '@babel/traverse';
 import { visitors } from '@babel/traverse';
-import type { AssignmentExpression, Identifier } from '@babel/types';
+import type { AssignmentExpression, Identifier, Property } from '@babel/types';
 import type { ComponentNode } from '../resolver';
 import type { Handler } from '.';
 
@@ -94,6 +96,121 @@ function findAssignedMethods(
   return state.methods;
 }
 
+// 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(
+  componentDefinition: NodePath<ComponentNode>,
+) {
+  let path: NodePath | null = componentDefinition;
+  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.
+    const value = resolveToValue(path);
+
+    if (value.isVariableDeclaration()) {
+      const decls: NodePath[] = value.get('declarations');
+
+      if (decls.length == 1) {
+        path = decls[0];
+      } else {
+        path = null;
+      }
+    } else if (value.isExpressionStatement()) {
+      path = value.get('expression');
+    } else if (value.isCallExpression()) {
+      if (isReactForwardRefCall(value) && !sawForwardRef) {
+        sawForwardRef = true;
+        path = value.get('arguments')[0];
+      } else {
+        path = null;
+      }
+    } else if (
+      value.isArrowFunctionExpression() ||
+      value.isFunctionDeclaration() ||
+      value.isFunctionExpression()
+    ) {
+      if (value.isArrowFunctionExpression()) {
+        path = value.get('body');
+      } else if (value.isFunctionDeclaration()) {
+        path = value.get('body');
+      } else if (value.isFunctionExpression()) {
+        path = value.get('body');
+      }
+
+      keepDigging = false;
+    } else {
+      // Any other type causes us to bail.
+      path = null;
+    }
+  }
+
+  return path;
+}
+
+function findImperativeHandleMethods(
+  componentDefinition: NodePath<ComponentNode>,
+): Array<NodePath<Property>> {
+  const path = findUnderlyingComponentDefinition(componentDefinition);
+
+  if (!path) {
+    return [];
+  }
+
+  const results: Array<NodePath<Property>> = [];
+
+  path.traverse({
+    CallExpression: function (callPath) {
+      // console.log('* call expression...');
+      // 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 (!arg.isArrowFunctionExpression()) {
+        return false;
+      }
+
+      const body = arg.get('body');
+
+      if (!body.isObjectExpression()) {
+        return false;
+      }
+
+      // We found the object body, now add all of the properties as methods.
+      body.get('properties').forEach(p => {
+        if (p.isObjectProperty()) {
+          results.push(p);
+        }
+      });
+
+      return false;
+    },
+  });
+
+  return results;
+}
+
 /**
  * Extract all flow types for the methods of a react component. Doesn't
  * return any react specific lifecycle methods.
@@ -103,7 +220,10 @@ const componentMethodsHandler: Handler = function (
   componentDefinition: NodePath<ComponentNode>,
 ): void {
   // Extract all methods from the class or object.
-  let methodPaths: Array<{ path: MethodNodePath; isStatic?: boolean }> = [];
+  let methodPaths: Array<{
+    path: MethodNodePath;
+    isStatic?: boolean;
+  }> = [];
 
   if (isReactComponentClass(componentDefinition)) {
     methodPaths = (
@@ -159,6 +279,15 @@ const componentMethodsHandler: Handler = function (
     ).map(p => ({ path: p }));
   }
 
+  // Also look for any methods that come from useImperativeHandle() calls.
+  const impMethodPaths = findImperativeHandleMethods(componentDefinition);
+
+  if (impMethodPaths && impMethodPaths.length > 0) {
+    methodPaths = methodPaths.concat(
+      impMethodPaths.map(p => ({ path: p as MethodNodePath })),
+    );
+  }
+
   documentation.set(
     'methods',
     methodPaths

packages/react-docgen/tests/integration/__fixtures__/component_44.tsx

@@ -0,0 +1,21 @@
+import React, { useRef, useImperativeHandle, forwardRef } from 'react';
+
+type Props = {
+  align?: "left" | "center" | "right" | "justify"
+};
+
+/**
+ * This is a TypeScript function component
+ */
+function FancyInput(props: Props, ref) {
+  const inputRef = useRef();
+  useImperativeHandle(ref, () => ({
+    /** this is a method on a component */
+    focus: () => {
+      inputRef.current.focus()
+    }
+  }));
+  return <input ref={inputRef} />;
+}
+
+export default forwardRef(FancyInput);