Add new resolver for finding annotated components (#743)

Author: danez

.changeset/stupid-files-buy.md

@@ -0,0 +1,13 @@
+---
+'react-docgen': minor
+---
+
+Added a new resolver that finds annotated components. This resolver is also
+enabled by default.
+
+To use this feature simply annotated a component with `@component`.
+
+```ts
+// @component
+class MyComponent {}
+```

.changeset/thin-foxes-talk.md

@@ -0,0 +1,8 @@
+---
+'react-docgen': major
+---
+
+Removed support for the `@extends React.Component` annotation on react class
+components.
+
+Instead you can use the new `@component` annotation.

packages/react-docgen/src/config.ts

@@ -16,7 +16,11 @@ import {
 import type { Importer } from './importer/index.js';
 import { fsImporter } from './importer/index.js';
 import type { Resolver } from './resolver/index.js';
-import { FindExportedDefinitionsResolver } from './resolver/index.js';
+import {
+  ChainResolver,
+  FindAnnotatedDefinitionsResolver,
+  FindExportedDefinitionsResolver,
+} from './resolver/index.js';
 
 export interface Config {
   handlers?: Handler[];
@@ -32,8 +36,14 @@ export interface Config {
 }
 export type InternalConfig = Omit<Required<Config>, 'filename'>;
 
-const defaultResolver: Resolver = new FindExportedDefinitionsResolver({
-  limit: 1,
+const defaultResolvers: Resolver[] = [
+  new FindExportedDefinitionsResolver({
+    limit: 1,
+  }),
+  new FindAnnotatedDefinitionsResolver(),
+];
+const defaultResolver: Resolver = new ChainResolver(defaultResolvers, {
+  chainingLogic: ChainResolver.Logic.ALL,
 });
 const defaultImporter: Importer = fsImporter;
 

packages/react-docgen/src/resolver/FindAnnotatedDefinitionsResolver.ts

@@ -0,0 +1,115 @@
+import normalizeClassDefinition from '../utils/normalizeClassDefinition.js';
+import type { NodePath } from '@babel/traverse';
+import { visitors } from '@babel/traverse';
+import type FileState from '../FileState.js';
+import type { ComponentNodePath, ResolverClass } from './index.js';
+import type {
+  ArrowFunctionExpression,
+  ClassDeclaration,
+  ClassExpression,
+  FunctionDeclaration,
+  FunctionExpression,
+  ObjectMethod,
+} from '@babel/types';
+
+interface TraverseState {
+  foundDefinitions: Set<ComponentNodePath>;
+  annotation: string;
+}
+
+function isAnnotated(path: NodePath, annotation: string): boolean {
+  let inspectPath: NodePath | null = path;
+
+  do {
+    const leadingComments = inspectPath.node.leadingComments;
+
+    // If an export doesn't have leading comments, we can simply continue
+    if (leadingComments && leadingComments.length > 0) {
+      // Search for the annotation in any comment.
+      const hasAnnotation = leadingComments.some(({ value }) =>
+        value.includes(annotation),
+      );
+
+      // if we found an annotation return true
+      if (hasAnnotation) {
+        return true;
+      }
+    }
+
+    // return false if the container of the current path is an array
+    // as we do not want to traverse up through this kind of nodes, like ArrayExpressions for example
+    // The only exception is variable declarations
+    if (
+      Array.isArray(inspectPath.container) &&
+      !inspectPath.isVariableDeclarator() &&
+      !inspectPath.parentPath?.isCallExpression()
+    ) {
+      return false;
+    }
+  } while ((inspectPath = inspectPath.parentPath));
+
+  return false;
+}
+
+function classVisitor(
+  path: NodePath<ClassDeclaration | ClassExpression>,
+  state: TraverseState,
+): void {
+  if (isAnnotated(path, state.annotation)) {
+    normalizeClassDefinition(path);
+    state.foundDefinitions.add(path);
+  }
+}
+
+function statelessVisitor(
+  path: NodePath<
+    | ArrowFunctionExpression
+    | FunctionDeclaration
+    | FunctionExpression
+    | ObjectMethod
+  >,
+  state: TraverseState,
+): void {
+  if (isAnnotated(path, state.annotation)) {
+    state.foundDefinitions.add(path);
+  }
+}
+
+const explodedVisitors = visitors.explode<TraverseState>({
+  ArrowFunctionExpression: { enter: statelessVisitor },
+  FunctionDeclaration: { enter: statelessVisitor },
+  FunctionExpression: { enter: statelessVisitor },
+  ObjectMethod: { enter: statelessVisitor },
+
+  ClassDeclaration: { enter: classVisitor },
+  ClassExpression: { enter: classVisitor },
+});
+
+interface FindAnnotatedDefinitionsResolverOptions {
+  annotation?: string;
+}
+
+/**
+ * Given an AST, this function tries to find all react components which
+ * are annotated with an annotation
+ */
+export default class FindAnnotatedDefinitionsResolver implements ResolverClass {
+  annotation: string;
+
+  constructor({
+    annotation = '@component',
+  }: FindAnnotatedDefinitionsResolverOptions = {}) {
+    this.annotation = annotation;
+  }
+
+  resolve(file: FileState): ComponentNodePath[] {
+    const state: TraverseState = {
+      foundDefinitions: new Set<ComponentNodePath>(),
+      annotation: this.annotation,
+    };
+
+    file.traverse(explodedVisitors, state);
+
+    return Array.from(state.foundDefinitions);
+  }
+}

packages/react-docgen/src/resolver/__tests__/FindAnnotatedDefinitionsResolver-test.ts

@@ -0,0 +1,217 @@
+import type { NodePath } from '@babel/traverse';
+import { parse, noopImporter } from '../../../tests/utils';
+import FindAnnotatedDefinitionsResolver from '../FindAnnotatedDefinitionsResolver.js';
+import { describe, expect, test } from 'vitest';
+
+describe('FindAnnotatedDefinitionsResolver', () => {
+  const resolver = new FindAnnotatedDefinitionsResolver();
+
+  function findComponentsInSource(
+    source: string,
+    importer = noopImporter,
+  ): NodePath[] {
+    return resolver.resolve(parse(source, {}, importer, true));
+  }
+
+  describe('class definitions', () => {
+    test('finds ClassDeclaration with line comment', () => {
+      const source = `
+        // @component
+        class Component {}
+      `;
+
+      expect(findComponentsInSource(source)).toMatchSnapshot();
+    });
+
+    test('finds ClassDeclaration with line comment and empty line', () => {
+      const source = `
+        // @component
+
+        class Component {}
+      `;
+
+      expect(findComponentsInSource(source)).toMatchSnapshot();
+    });
+
+    test('finds ClassDeclaration with block comment', () => {
+      const source = `
+        /* @component */
+        class Component {}
+      `;
+
+      expect(findComponentsInSource(source)).toMatchSnapshot();
+    });
+
+    test('finds ClassDeclaration with block comment and empty line', () => {
+      const source = `
+        /* @component */
+
+        class Component {}
+      `;
+
+      expect(findComponentsInSource(source)).toMatchSnapshot();
+    });
+
+    test('finds ClassExpression', () => {
+      const source = `
+        // @component
+        const Component = class {}
+      `;
+
+      expect(findComponentsInSource(source)).toMatchSnapshot();
+    });
+
+    test('finds ClassExpression with assignment', () => {
+      const source = `
+        let Component;
+        // @component
+        Component = class {}
+      `;
+
+      expect(findComponentsInSource(source)).toMatchSnapshot();
+    });
+
+    test('finds nothing when not annotated', () => {
+      const source = `
+        class ComponentA {}
+        const ComponentB = class {}
+      `;
+
+      const result = findComponentsInSource(source);
+
+      expect(Array.isArray(result)).toBe(true);
+      expect(result.length).toBe(0);
+    });
+  });
+
+  describe('stateless components', () => {
+    test('finds ArrowFunctionExpression', () => {
+      const source = `
+        // @component
+        const Component = () => {};
+      `;
+
+      expect(findComponentsInSource(source)).toMatchSnapshot();
+    });
+
+    test('finds FunctionDeclaration', () => {
+      const source = `
+        // @component
+        function Component() {}
+      `;
+
+      expect(findComponentsInSource(source)).toMatchSnapshot();
+    });
+
+    test('finds FunctionExpression', () => {
+      const source = `
+        // @component
+        const Component = function() {};
+      `;
+
+      expect(findComponentsInSource(source)).toMatchSnapshot();
+    });
+
+    test('finds ObjectMethod', () => {
+      const source = `
+        const obj = {
+          // @component
+          component() {}
+        };
+      `;
+
+      expect(findComponentsInSource(source)).toMatchSnapshot();
+    });
+
+    test('finds ObjectProperty', () => {
+      const source = `
+        const obj = {
+          // @component
+          component: function() {}
+        };
+      `;
+
+      expect(findComponentsInSource(source)).toMatchSnapshot();
+    });
+
+    test('finds nothing when not annotated', () => {
+      const source = `
+        const ComponentA = () => {};
+        function ComponentB() {}
+        const ComponentC = function() {};
+        const obj = { component() {} };
+        const obj2 = { component: function() {} };
+      `;
+
+      const result = findComponentsInSource(source);
+
+      expect(Array.isArray(result)).toBe(true);
+      expect(result.length).toBe(0);
+    });
+  });
+
+  test('finds component wrapped in HOC', () => {
+    const source = `
+        // @component
+        const Component = React.memo(() => {});
+      `;
+
+    expect(findComponentsInSource(source)).toMatchSnapshot();
+  });
+
+  test('finds component wrapped in two HOCs', () => {
+    const source = `
+        // @component
+        const Component = React.memo(otherHOC(() => {}));
+      `;
+
+    expect(findComponentsInSource(source)).toMatchSnapshot();
+  });
+
+  test('finds component wrapped in function', () => {
+    const source = `
+        function X () {
+          // @component
+          const subcomponent = class {}
+        }
+      `;
+
+    expect(findComponentsInSource(source)).toMatchSnapshot();
+  });
+
+  test('does not traverse up ArrayExpressions', () => {
+    const source = `
+        // @component
+        const arr = [
+          function() {},
+          function() {}
+        ]
+      `;
+
+    const result = findComponentsInSource(source);
+
+    expect(Array.isArray(result)).toBe(true);
+    expect(result.length).toBe(0);
+  });
+
+  test('does not traverse up function parameters', () => {
+    const source = `
+        // @component
+        function x(component = () => {}) {}
+      `;
+
+    const result = findComponentsInSource(source);
+
+    expect(Array.isArray(result)).toBe(true);
+    expect(result.length).toBe(1);
+    expect(result[0].type).not.toBe('ArrowFunctionExpression');
+  });
+
+  test('finds function parameter with annotation', () => {
+    const source = `
+        function x(component = /* @component */() => {}) {}
+      `;
+
+    expect(findComponentsInSource(source)).toMatchSnapshot();
+  });
+});