fix: infer displayName given forwardRef() node

motiz88 · Feb 13, 2020 · 1977147 · 1977147
1 parent d8002c9
commit 1977147
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 1 deletion.
diff --git a/src/handlers/__tests__/displayNameHandler-test.js b/src/handlers/__tests__/displayNameHandler-test.js
@@ -181,6 +181,17 @@ describe('defaultPropsHandler', () => {
       expect(documentation.displayName).toBe('Foo');
     });
 
+    it('considers the variable name when handling forwardRef', () => {
+      const definition = statement(
+        [
+          'var Foo = React.forwardRef(() => {});',
+          'import React from "react";',
+        ].join('\n'),
+      ).get('declarations', 0, 'init');
+      expect(() => displayNameHandler(documentation, definition)).not.toThrow();
+      expect(documentation.displayName).toBe('Foo');
+    });
+
     it('considers the variable name on assign', () => {
       const definition = statement('Foo = () => {};').get(
         'expression',
@@ -201,6 +212,18 @@ describe('defaultPropsHandler', () => {
       expect(documentation.displayName).toBe('Foo');
     });
 
+    it('considers the variable name on assign when handling forwardRef call', () => {
+      const definition = statement([
+        'Foo = React.forwardRef(() => {});',
+        'import React from "react";',
+      ].join('\n'),).get(
+        'expression',
+        'right',
+      );
+      expect(() => displayNameHandler(documentation, definition)).not.toThrow();
+      expect(documentation.displayName).toBe('Foo');
+    });
+
     it('considers a static displayName object property over variable name', () => {
       const definition = statement(`
         var Foo = () => {};

diff --git a/src/handlers/displayNameHandler.js b/src/handlers/displayNameHandler.js
@@ -10,6 +10,7 @@
 import { namedTypes as t } from '@motiz88/ast-types';
 import getMemberValuePath from '../utils/getMemberValuePath';
 import getNameOrValue from '../utils/getNameOrValue';
+import isReactForwardRefCall from '../utils/isReactForwardRefCall';
 import resolveToValue from '../utils/resolveToValue';
 import resolveFunctionDefinitionToReturnValue from '../utils/resolveFunctionDefinitionToReturnValue';
 import type Documentation from '../Documentation';
@@ -29,7 +30,8 @@ export default function displayNameHandler(
       documentation.set('displayName', getNameOrValue(path.get('id')));
     } else if (
       t.ArrowFunctionExpression.check(path.node) ||
-      t.FunctionExpression.check(path.node)
+      t.FunctionExpression.check(path.node) ||
+      isReactForwardRefCall(path)
     ) {
       let currentPath = path;
       while (currentPath.parent) {