codemod-forwardRef.js

// Press ctrl+space for code completion
export default function transformer(file, api) {
  const j = api.jscodeshift;

  function hasPropsArgument(node) {
    return node.params.length === 0 || node.params[0].name === 'props';
  }

  function isFunctionComponent(path) {
    const { node } = path;
    return /^[A-Z]/.test(node.id.name) && hasPropsArgument(node);
  }

  function forwardRefCall(name, argument) {
    return j.variableDeclaration('const', [
      j.variableDeclarator(
        j.identifier(name),
        j.callExpression(j.memberExpression(j.identifier('React'), j.identifier('forwardRef')), [
          argument,
        ]),
      ),
    ]);
  }

  function wrapFunctionComponent(path) {
    const {
      node: {
        body,
        id,
        params: [propsParam],
      },
    } = path;
    const params = [propsParam, j.identifier('ref')];

    return forwardRefCall(id.name, j.functionExpression(id, params, body));
  }

  function wrapFunctionComponents(program) {
    return (
      program
        .find(j.FunctionDeclaration)
        .filter(isFunctionComponent)
        .replaceWith(wrapFunctionComponent).length > 0
    );
  }

  function wrapClassComponent(path) {
    const { node } = path;

    return j.callExpression(node.callee, [
      j.callExpression(j.identifier('withForwardRef'), node.arguments),
    ]);
  }

  function createPropType(name) {
    return j.memberExpression(j.identifier('PropTypes'), j.identifier(name));
  }

  function refPropType() {
    return j.callExpression(createPropType('oneOfType'), [
      j.arrayExpression([createPropType('func'), createPropType('object')]),
    ]);
  }

  function addInnerRefPropType(path) {
    const { node } = path;
    return j.assignmentExpression(
      node.operator,
      node.left,
      j.objectExpression([
        ...node.right.properties,
        j.objectProperty(j.identifier('innerRef'), refPropType()),
      ]),
    );
  }

  /**
   * adds an innerRef property to a propTypes object if a withStyles hoc is used
   * @param {*} program
   */
  function wrapClassComponents(program) {
    // withStyles(styles)(Componnet) => withStyles(styles)(withForwardRef(Component))
    const classComponentDecorations = program.find(j.CallExpression, {
      callee: {
        callee: {
          type: 'Identifier',
          name: 'withStyles',
        },
      },
    });

    // use withForwardRefHoC
    classComponentDecorations.replaceWith(wrapClassComponent);

    const usedHoc = classComponentDecorations.length > 0;
    if (usedHoc) {
      // add innerRef propTypes
      program
        .find(j.AssignmentExpression, {
          left: {
            type: 'MemberExpression',
            property: {
              name: 'propTypes',
            },
          },
        })
        .replaceWith(addInnerRefPropType);

      return true;
    }

    return false;
  }

  const program = j(file.source);

  const didContainFunctionComponent = wrapFunctionComponents(program);
  if (!didContainFunctionComponent) {
    wrapClassComponents(program);
  }

  return program.toSource();
}